mstar.engine.resources

Contents

mstar.engine.resources#

The resource layer’s declaration surface.

Everything re-exported here is a declaration — a spec, a per-request config, a step, or the generic step envelope — and every module behind it is free of the managers and their kernels. So a model can import all of it in one line without dragging FlashInfer or Triton in behind it; the concrete resources are reached by their own paths.

Nothing inside the package may import from here: during this module’s own execution the package is only half-initialized. Import siblings by their submodule path instead.

class mstar.engine.resources.AdmitFailedReason(message: str)[source]#

Bases: object

Parameters:

message (str)

message: str#
terminal: ClassVar[bool] = False#
class mstar.engine.resources.AdmitOutcome(ok, ready, reason)[source]#

Bases: NamedTuple

Parameters:
ok: bool#

Alias for field number 0

ready: bool#

Alias for field number 1

reason: AdmitFailedReason | None#

Alias for field number 2

class mstar.engine.resources.AdmitRuntimeError(message, request_ids=frozenset({}))[source]#

Bases: AdmitFailedReason

A resource cannot serve this request at all.

Terminal, unlike the two above: no eviction and no reload makes it go away, so the caller’s answer is to fail the request, not to retry it. request_ids names the requests at fault; empty means the refusal is about the step as a whole.

Parameters:
request_ids: frozenset[str] = frozenset({})#
terminal: ClassVar[bool] = True#
class mstar.engine.resources.AllocationFailed(message: str, pages_short: int, label: str, request_id: str, rows_fit: int | None = None, rows_fit_alone: int | None = None, rows_sized: int | None = None)[source]#

Bases: AdmitFailedReason

Parameters:

  • message (str)

  • pages_short (int)

  • label (str)

  • request_id (str)

  • rows_fit (int | None)

  • rows_fit_alone (int | None)

  • rows_sized (int | None)

label: str#
pages_short: int#
request_id: str#
rows_fit: int | None = None#
rows_fit_alone: int | None = None#
rows_sized: int | None = None#
class mstar.engine.resources.AttentionConfig(kv_cache: str, backend: mstar.engine.resources.attn.config.AttnBackend = <AttnBackend.AUTO: 'auto'>, flashinfer_backend: str = 'auto', sm_scale: float | None = None, splits: tuple[mstar.engine.resources.spec.CombinedSplit, ...]=(), sliding_window: int | None = None)[source]#

Bases: object

Parameters:
backend: AttnBackend = 'auto'#
flashinfer_backend: str = 'auto'#
kv_cache: str#
sliding_window: int | None = None#
sm_scale: float | None = None#
splits: tuple[CombinedSplit, ...] = ()#
class mstar.engine.resources.AttentionSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.attn.config.AttentionConfig)[source]#

Bases: NodeResourceSpec

Parameters:
apply_yaml_overrides(backend=None, flashinfer_backend=None)[source]#

Which kernel to run is the deployment’s call as much as the model’s — an image that cannot build FA3 pins FA2 here.

Cache geometry is not repeated here: it belongs to the KV resource this spec depends on, and is tuned under that resource’s own block.

Parameters:
config: AttentionConfig#
depends_on()[source]#

Keys whose specs this one builds against; the engine resolves them into EngineResourceInfo.dependencies.

Return type:

set[str]

property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

class mstar.engine.resources.AttentionStep(segments: tuple[mstar.engine.resources.step.Segment, ...] | None = None, causal: bool = True, context_only: bool = False)[source]#

Bases: ResourceStep

Parameters:
causal: bool = True#
context_only: bool = False#
class mstar.engine.resources.AttnBackend(*values)[source]#

Bases: Enum

AUTO = 'auto'#
DENSE = 'dense'#
FLASHINFER = 'flashinfer'#
FLASHINFER_MLA = 'flashinfer_mla'#
FLEX = 'flex'#
XPU_PAGED = 'xpu_paged'#
class mstar.engine.resources.BucketKey(graph_walk: str, bs: int, num_tokens: int, cg_key_info: Any | None = None)[source]#

Bases: object

Parameters:

  • graph_walk (str)

  • bs (int)

  • num_tokens (int)

  • cg_key_info (Any | None)

bs: int#
cg_key_info: Any | None = None#
graph_walk: str#
num_tokens: int#
class mstar.engine.resources.CGSlotSpec(bucket: 'BucketKey', slot: 'int', config: 'AcceleratorGraphConfig | PiecewiseAcceleratorGraphConfig', config_idx: 'int | None' = None)[source]#

Bases: object

Parameters:
property bs#
bucket: BucketKey#
config: AcceleratorGraphConfig | PiecewiseAcceleratorGraphConfig#
config_idx: int | None = None#
property num_tokens#
slot: int#
class mstar.engine.resources.CrossAttentionConfig(kv_cache, query_kv_cache=None, context_label='context', backend=AttnBackend.AUTO, flashinfer_backend='auto')[source]#

Bases: object

Cross-attention against a context written once and never extended.

kv_cache names the KV resource holding the encoder context; query_kv_cache names the decoder’s KV resource, whose plan defines this step’s query packing. They may be the same resource when the context shares the decoder’s head config — the context then lives in it under its own context_label. They differ when it does not, which is the usual case (an encoder’s head count rarely matches the decoder’s).

query_kv_cache=None covers the query side having no KV cache at all (nothing is cached across steps on it): the packing then comes off the cross-attention step’s own segments, one qo entry per segment in declared order.

Parameters:
backend: AttnBackend = 'auto'#
context_label: str = 'context'#
flashinfer_backend: str = 'auto'#
kv_cache: str#
query_kv_cache: str | None = None#
class mstar.engine.resources.CrossAttentionSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.attn.config.CrossAttentionConfig)[source]#

Bases: NodeResourceSpec

Parameters:
apply_yaml_overrides(backend=None, flashinfer_backend=None)[source]#

Which kernel to run against the context cache; see AttentionSpec.

Parameters:
config: CrossAttentionConfig#
depends_on()[source]#

Keys whose specs this one builds against; the engine resolves them into EngineResourceInfo.dependencies.

Return type:

set[str]

property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

class mstar.engine.resources.FullAdmitOutcome(outcome, failed_resource=None)[source]#

Bases: NamedTuple

What the runner answers with: one resource’s outcome, plus which resource gave it.

A resource doesn’t know the key it is registered under, so the runner — which does — names it on the way out. The caller needs it to scope an eviction to the resource that actually ran out.

Parameters:
failed_resource: str | None#

Alias for field number 1

property ok: bool#
outcome: AdmitOutcome#

Alias for field number 0

property ready: bool#
property reason: AdmitFailedReason | None#
class mstar.engine.resources.KVConfig(*, num_layers, num_kv_heads, head_dim, num_qo_heads=None)[source]#

Bases: ABC

The model geometry every KV storage strategy needs, and nothing else.

What a cache holds (layers, heads, head dim) is a checkpoint fact and lives here; how it is stored — paged or ringed, is storage policy and lives in a subclass. KVSpec dispatches on which one it was handed.

kw_only is load-bearing, not style: num_qo_heads carries a default here while both subclasses add required fields, which is an invalid positional dataclass field order.

Parameters:

  • num_layers (int)

  • num_kv_heads (int)

  • head_dim (int)

  • num_qo_heads (int | None)

abstractmethod apply_yaml_overrides(**kwargs)[source]#

Patch this deployment’s tunables; see NodeResourceSpec.

Abstract so each storage policy names exactly what it accepts and lets the rest raise — a paged key against a ring config is a typo, and the spec’s contract is that a typo is loud.

Return type:

None

head_dim: int#
num_kv_heads: int#
num_layers: int#
num_qo_heads: int | None = None#
shard(num_shards)[source]#

Narrow the head counts to one rank’s slice.

Idempotent because one KVConfig is shared by the KV resource and the attention resources planned against it, and each shards on construction. num_shards is the instance session size (tp * sp): Ulysses SP all-to-alls heads, so attention runs at head-degree tp*sp.

Parameters:

num_shards (int)

Return type:

None

class mstar.engine.resources.KVLayout(*values)[source]#

Bases: Enum

MLA = 'MLA'#
NHD = 'NHD'#
class mstar.engine.resources.KVReqConfig(needed_labels: list[str] | None = None, needed_labels_per_node: dict[str, list[str]] = <factory>, needed_labels_per_node_walk: dict[tuple[str, str], list[str]] = <factory>, publish_labels_per_node_walk: dict[tuple[str, str], list[str]] | None = None, final_publish_labels_per_node_walk: dict[tuple[str, str], list[str]] | None = None, prefix_keys: dict[str, list[bytes]] | None = None, prefix_tail: dict[str, list[int]] | None = None, prefix_decode: dict[str, str] | None = None, prefix_cache: bool = True, kv_layout: dict[str, 'StreamLayout'] | None = None, loop_caps: dict[str, int] | None = None)[source]#

Bases: ResourceReqConfig

Parameters:
apply_conductor_config(prefix_keys=None, prefix_tail=None, prefix_decode=None, prefix_cache=None, kv_layout=None, loop_caps=None, **kwargs)[source]#
Parameters:
final_publish_labels_per_node_walk: dict[tuple[str, str], list[str]] | None = None#
get_labels(node, walk)[source]#
Parameters:
get_publish_labels(node, walk, available, *, final=False)[source]#
Parameters:
Return type:

list[str]

kv_layout: dict[str, StreamLayout] | None = None#
loop_caps: dict[str, int] | None = None#
loop_room(spec=None)[source]#

The most iterations each loop can run before the request outgrows the resource built from spec. A loop the resource doesn’t bound is left out.

Parameters:

spec (KVSpec | None)

Return type:

dict[str, int]

needed_labels: list[str] | None = None#
needed_labels_per_node: dict[str, list[str]]#
needed_labels_per_node_walk: dict[tuple[str, str], list[str]]#
opened_labels()[source]#
Return type:

set[str]

prefix_cache: bool = True#
prefix_decode: dict[str, str] | None = None#
prefix_keys: dict[str, list[bytes]] | None = None#
prefix_tail: dict[str, list[int]] | None = None#
publish_labels_per_node_walk: dict[tuple[str, str], list[str]] | None = None#
validate(spec=None)[source]#

ValueError (a 400) when a label would hold more positions at the request’s caps than its pool’s max_seq_len. A label under retention is left out, since Cosmos3, the one model that declares one, computes its positions in each forward.

Parameters:

spec (KVSpec | None)

Return type:

None

class mstar.engine.resources.KVSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.kv.config.KVConfig)[source]#

Bases: NodeResourceSpec

Parameters:
apply_yaml_overrides(**kwargs)[source]#

Forwarded to the config: which keys are legal is a property of the storage strategy, so the config answers for them.

config: KVConfig#
property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

class mstar.engine.resources.KVStep(segments: tuple[mstar.engine.resources.step.Segment, ...] | None = None, commit: bool = True, combined_labels: dict[tuple[str, ...], str] = <factory>, pre_forks: tuple[tuple[str, str], ...] = (), post_forks: tuple[tuple[str, str], ...] = (), fork_rids: frozenset | None = None, sink: frozenset[tuple[str, str]] = frozenset(), retention: collections.abc.Mapping[tuple[str, str], mstar.engine.resources.kv.config.RetentionPolicy] = <factory>)[source]#

Bases: ResourceStep

Parameters:
combined_labels: dict[tuple[str, ...], str]#
commit: bool = True#
fork_rids: frozenset | None = None#
post_forks: tuple[tuple[str, str], ...] = ()#
pre_forks: tuple[tuple[str, str], ...] = ()#
retention: Mapping[tuple[str, str], RetentionPolicy]#
sink: frozenset[tuple[str, str]] = frozenset({})#
class mstar.engine.resources.NodeResourceSpec(resource_key: str, nodes: set[str])[source]#

Bases: ABC

Parameters:
apply_yaml_overrides(**kwargs)[source]#

Patch declared parameters from this resource’s YAML block.

The model declares shapes that suit the model; a deployment tunes what suits the box it runs on. The block is scoped to this spec’s resource_key (see apply_yaml_overrides below), so an unrecognized key here is a typo, not another resource’s setting — subclasses name exactly what they accept and let the rest raise.

TODO: generalize. Matching YAML keys against the spec’s (and its config’s) dataclass fields would remove these per-resource implementations, at the cost of silently accepting anything named alike.

depends_on()[source]#

Keys whose specs this one builds against; the engine resolves them into EngineResourceInfo.dependencies.

Return type:

set[str]

nodes: set[str]#
abstract property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

resource_key: str#
class mstar.engine.resources.PagedKVConfig(*, num_layers, num_kv_heads, head_dim, num_qo_heads=None, max_seq_len, max_num_pages=2048, page_size=128, layout=KVLayout.NHD, cpu_offload_pages=0, prefix_cache_salt='', prefix_cache=True, prefix_cache_tp=False, kv_lora_rank=None, qk_rope_head_dim=None)[source]#

Bases: KVConfig

Fixed-size pages, appended to as a sequence grows. The default.

Parameters:

  • num_layers (int)

  • num_kv_heads (int)

  • head_dim (int)

  • num_qo_heads (int | None)

  • max_seq_len (int)

  • max_num_pages (int)

  • page_size (int)

  • layout (KVLayout)

  • cpu_offload_pages (int)

  • prefix_cache_salt (str)

  • prefix_cache (bool)

  • prefix_cache_tp (bool)

  • kv_lora_rank (int | None)

  • qk_rope_head_dim (int | None)

apply_yaml_overrides(max_num_pages=None, page_size=None, max_seq_len=None, cpu_offload_pages=None, prefix_cache_salt=None, prefix_cache=None, prefix_cache_tp=None)[source]#

How much cache this deployment gets, and how it is cut up.

Parameters:

  • max_num_pages (int | None)

  • page_size (int | None)

  • max_seq_len (int | None)

  • cpu_offload_pages (int | None)

  • prefix_cache_salt (str | None)

  • prefix_cache (bool | None)

  • prefix_cache_tp (bool | None)

Return type:

None

cpu_offload_pages: int = 0#
kv_lora_rank: int | None = None#
property latent_dim: int#

Per-token cache width under MLA.

layout: KVLayout = 'NHD'#
max_num_pages: int = 2048#
max_seq_len: int#
page_size: int = 128#
prefix_cache: bool = True#
prefix_cache_salt: str = ''#
prefix_cache_tp: bool = False#
qk_rope_head_dim: int | None = None#
qo_heads_given: bool = False#
class mstar.engine.resources.PosBackend(*values)[source]#

Bases: Enum

ROPE = 'rope'#
class mstar.engine.resources.PosScheme(*values)[source]#

Bases: Enum

BLOCK = 'block'#
SEQUENTIAL = 'sequential'#
class mstar.engine.resources.PositionConfig(kv_cache: str, backend: mstar.engine.resources.position.config.PosBackend = <PosBackend.ROPE: 'rope'>, scheme: mstar.engine.resources.position.config.PosScheme = <PosScheme.SEQUENTIAL: 'sequential'>, block_step: int = 1, rotary_dim: int | None = None, interleave: bool = False, rope_scale: float = 1.0, rope_theta: float = 10000.0, rope_dtype: torch.dtype | None = None, rope_table_dtype: torch.dtype | None = None, low_freq_factor: float | None = None, high_freq_factor: float | None = None, old_context_len: int | None = None)[source]#

Bases: object

Parameters:
backend: PosBackend = 'rope'#
block_step: int = 1#
high_freq_factor: float | None = None#
interleave: bool = False#
kv_cache: str#
property llama31_params: dict[str, float]#
low_freq_factor: float | None = None#
old_context_len: int | None = None#
rope_dtype: dtype | None = None#
rope_scale: float = 1.0#
rope_table_dtype: dtype | None = None#
rope_theta: float = 10000.0#
rotary_dim: int | None = None#
scheme: PosScheme = 'sequential'#
class mstar.engine.resources.PositionSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.position.config.PositionConfig)[source]#

Bases: NodeResourceSpec

Parameters:
config: PositionConfig#
depends_on()[source]#

Keys whose specs this one builds against; the engine resolves them into EngineResourceInfo.dependencies.

Return type:

set[str]

property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

class mstar.engine.resources.PositionStep(segments: tuple[mstar.engine.resources.step.Segment, ...] | None = None, pos_ids: 'dict[str, torch.Tensor] | torch.Tensor | None' = None, advance: tuple[int, ...] | None = None)[source]#

Bases: ResourceStep

Parameters:
advance: tuple[int, ...] | None = None#
pos_ids: dict[str, Tensor] | Tensor | None = None#
class mstar.engine.resources.PublishedInfo[source]#

Bases: ABC

abstractmethod clone()[source]#

Return independent publication metadata for deferred sending.

Return type:

PublishedInfo

abstractmethod update(other)[source]#
Parameters:

other (PublishedInfo)

Return type:

None

class mstar.engine.resources.RaggedAttentionConfig(num_qo_heads, num_kv_heads, head_dim, sm_scale=None, max_segments_per_request=1, max_tokens_per_request=None, backend=RaggedAttnBackend.AUTO, flashinfer_backend='auto', dtype=None)[source]#

Bases: object

Varlen attention over segments packed into one forward, with no KV cache: the whole layout is this step’s, and nothing carries to the next. The config of both RaggedAttentionSpec (self-attention within each span) and RaggedCrossAttentionSpec (one span attending another).

Head counts are pre-sharding; the engine narrows them to the rank’s slice at build, as it does for a KVConfig.

Parameters:
backend: RaggedAttnBackend = 'auto'#
dtype: dtype | None = None#
flashinfer_backend: str = 'auto'#
head_dim: int#
max_segments_for(bs)[source]#
Parameters:

bs (int)

Return type:

int

max_segments_per_request: int = 1#
max_tokens_for(bs)[source]#
Parameters:

bs (int)

Return type:

int | None

max_tokens_per_request: int | None = None#
num_kv_heads: int#
num_qo_heads: int#
shard(num_shards)[source]#

Narrow the head counts to one rank’s slice; see KVConfig.shard.

Idempotent, so a rebuild (or a second manager over one config) is free.

Parameters:

num_shards (int)

Return type:

None

sm_scale: float | None = None#
class mstar.engine.resources.RaggedAttentionSpec(resource_key, nodes, config)[source]#

Bases: _RaggedSpecBase

Cacheless self-attention within each declared span (AttentionStep).

Parameters:
property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

class mstar.engine.resources.RaggedAttnBackend(*values)[source]#

Bases: Enum

AUTO = 'auto'#
FLASHINFER = 'flashinfer'#
class mstar.engine.resources.RaggedBlockCausalAttentionSpec(resource_key, nodes, config, block_size)[source]#

Bases: _RaggedSpecBase

Cacheless block-causal self-attention within each declared span: a token in block b of its span (block_size tokens a block, from the span’s start) attends every key of blocks 0..b. Declared with AttentionStep(causal=False), since the block structure is the whole mask.

Its own kind, not a flag on RaggedAttentionSpec: it runs a different kernel (paged prefill over one-token pages, see attn.ragged.block_causal).

Parameters:
block_size: int#
property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

class mstar.engine.resources.RaggedCrossAttentionSpec(resource_key, nodes, config)[source]#

Bases: _RaggedSpecBase

Cacheless cross-attention between two spans of each request: a forward’s (q_label, kv_label) pairs (RaggedCrossAttentionStep), each run through RaggedAttentionCallable(resource, cross_label(q_label, kv_label)).

Separate from RaggedAttentionSpec so a resource is one kind of attention: a model with both declares one of each per head geometry.

Parameters:
property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

class mstar.engine.resources.RaggedCrossAttentionStep(segments=None, pairs=())[source]#

Bases: ResourceStep

A ragged cross-attention step: which spans attend which.

Each label in segments is one span per request (video tokens, text tokens, …). Each (q_label, kv_label) in pairs is one attention: every request’s q_label span attends that request’s kv_label span, so both labels must be declared for the same requests in the same order. Never causal.

Parameters:
pairs: tuple[tuple[str, str], ...] = ()#
class mstar.engine.resources.RequestOffloading(message, label, request_id)[source]#

Bases: AdmitFailedReason

The request’s state is moving to host memory; retry once it is back.

Distinct from AllocationFailed because the answer is different: nothing needs evicting, the caller just re-drives the step once reload has run.

Parameters:
label: str#
request_id: str#
class mstar.engine.resources.Resource[source]#

Bases: ABC

admit(step, ctx)[source]#

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.

Parameters:
Return type:

AdmitOutcome

admit_retrieve(rid, node_name, graph_walk, published)[source]#

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.

Parameters:
Return type:

AdmitOutcome

adopt_cached_prefix(rid, node_name, graph_walk, pages, caught_up)[source]#

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.

Parameters:
Return type:

bool

apply_cached_prefix(rid, node_name, graph_walk, inputs, prefix)[source]#

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.

Parameters:
Return type:

None

abstractmethod classmethod build(spec, info)[source]#
Parameters:
Return type:

Resource

build_accelerator_graph_buffers(slots, max_bs, max_seq_len)[source]#

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.

Parameters:
Return type:

None

can_capture()[source]#
Return type:

bool

can_reload(rid)[source]#

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.

Parameters:

rid (str)

Return type:

bool

cleanup()[source]#
clear_preplan()[source]#
commit(step, ctx)[source]#

record step consumption

Parameters:
Return type:

None

complete_cached_walk(rid, node_name, graph_walk)[source]#

Called for a walk the cache served whole, since that walk never admits or commits.

Parameters:
Return type:

None

depends_on()[source]#
Return type:

set[str]

enable_prefix_cache(root, walks=None)[source]#

Open whatever this resource keeps across requests, under root; True if anything opened.

Parameters:
Return type:

bool

extend_prefix_chain(rid, node_name, graph_walk, outputs)[source]#

Take this step’s sampled tokens, so what was generated can be keyed.

Parameters:
Return type:

None

fingerprint()[source]#

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 type:

bytes | None

property force_double_buffer#

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.

get_offload_priority(rid)[source]#

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.

Parameters:

rid (str)

Return type:

float

ingest_capture_row(rid)[source]#

Open a CUDA graph capture padding row; by default, as a request.

Parameters:

rid (str)

ingest_request(rid, overrides)[source]#

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.

Parameters:
is_offloaded(rid)[source]#
Parameters:

rid (str)

Return type:

bool

property keys_prefix_chains: 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.

lease_cached_prefix(rid, node_name, graph_walk)[source]#

On rank 0 of a TP node, lease the request’s cached prefix as its step is sent and return the pages, or None.

Parameters:
Return type:

list[int] | None

offload(rid)[source]#

Move the request’s state off-device. Returns what was reclaimed.

Parameters:

rid (str)

Return type:

int

page_trail()[source]#

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 type:

tuple[int, int] | None

plan(step, ctx)[source]#

ret is immutable and opaque to runner; only gives to ctx.plan_results

Parameters:
Return type:

Any

post_warmup_validate()[source]#

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.

prefix_skip_safe: bool = False#
publish(request_id)[source]#
Parameters:

request_id (str)

Return type:

PublishedInfo | None

publish_after_stop(request_id, node_name, graph_walk)[source]#

Publish state that is useful only when a dynamic loop stops.

Parameters:

  • request_id (str)

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

PublishedInfo | None

publish_for_step(request_id, node_name, graph_walk)[source]#

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.

Parameters:

  • request_id (str)

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

PublishedInfo | None

publish_from_snapshot(request_id, snapshot, node_name, graph_walk)[source]#

Finish a publish_snapshot_for_step record into the published info.

Parameters:

  • request_id (str)

  • snapshot (Any)

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

PublishedInfo | None

publish_snapshot_batch(request_ids, node_name, graph_walk)[source]#

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.

Parameters:

  • request_ids (list[str])

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

list

publish_snapshot_for_step(request_id, node_name, graph_walk)[source]#

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.

Parameters:

  • request_id (str)

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

Any

reclaimable(rid)[source]#

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.

Parameters:

rid (str)

Return type:

int

release_cached_prefix(rid, node_name, graph_walk)[source]#

Give back what lease_cached_prefix or adopt_cached_prefix took, for a step that will not run.

Parameters:
Return type:

None

reload(rid)[source]#

Bring it back. False when it doesn’t fit on device yet.

Parameters:

rid (str)

Return type:

bool

remove_request(rid)[source]#
Parameters:

rid (str)

reset_request(rid, free=False)[source]#

For clearing dummy RIDs during cuda graph capture

Parameters:
resolve_cached_prefix(rid, node_name, graph_walk, inputs)[source]#

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.

Parameters:
Return type:

CachedPrefix | None

rollback_admit(step, ctx)[source]#

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.

Parameters:
Return type:

None

property supports_eviction#
property supports_preplan#
class mstar.engine.resources.ResourceReqConfig[source]#

Bases: object

Per-request parameters for one resource, carried on the request and handed to that resource at ingest. Keyed by resource key, so it needs no tag of its own — a marker base, with no contract beyond the hooks below.

apply_conductor_config(**kwargs)[source]#
loop_room(spec=None)[source]#

The most iterations each loop can run before the request outgrows the resource built from spec. A loop the resource doesn’t bound is left out.

Parameters:

spec (NodeResourceSpec | None)

Return type:

dict[str, int]

validate(spec=None)[source]#

Raise ValueError if a client-supplied value is out of range, or one the resource built from spec cannot honour.

The conductor calls it at admission, so a bad request is a 400 before any worker sees it.

Parameters:

spec (NodeResourceSpec | None)

Return type:

None

class mstar.engine.resources.ResourceStep(segments: tuple[mstar.engine.resources.step.Segment, ...] | None = None)[source]#

Bases: object

Parameters:

segments (tuple[Segment, ...] | None)

segments: tuple[Segment, ...] | None = None#
class mstar.engine.resources.RingKVConfig(*, num_layers, num_kv_heads, head_dim, num_qo_heads=None, tokens_per_frame, layers, num_sessions=1)[source]#

Bases: KVConfig

A fixed horizon of frame slots per layer, overwritten in place.

Every layer’s storage is allocated once and reused for the life of the process, and a write to an occupied slot is the intended behaviour.

Parameters:
apply_yaml_overrides(num_sessions=None, **kwargs)[source]#

num_sessions only. Nothing else here is a deployment knob.

Parameters:

num_sessions (int | None)

Return type:

None

layers: tuple[RingKVLayerConfig, ...]#
num_sessions: int = 1#
tokens_per_frame: int#
property total_sessions: int#

Resident sessions plus one shared scratch session for padding rows.

A replay padded to its capture bucket parks the dummy tail on this session (see RingKVManager.plan); it is never handed to a request, so resident capacity stays num_sessions and the deployment knob keeps its meaning. The ring buffer and the flex mask both size on this count so the padding session is a real, addressable span.

class mstar.engine.resources.RingKVLayerConfig(ring_frames, ring_buckets, pinned_dilation)[source]#

Bases: object

One layer’s ring geometry. Layers can hold frames at different strides.

Parameters:

  • ring_frames (int)

  • ring_buckets (int)

  • pinned_dilation (int)

pinned_dilation: int#
ring_buckets: int#
ring_frames: int#
class mstar.engine.resources.RingKVStep(segments=None, *, frames)[source]#

Bases: ResourceStep

One ring clock per request in the step’s batch.

Parameters:
frames: tuple[tuple[str, int], ...]#
class mstar.engine.resources.SamplerSpec(resource_key: str, nodes: set[str], vocab_size: int | None, enable_repetion_penalty: bool = True, enable_min_p: bool = False, max_repetition_window: int = 0, min_tokens_stop_ids: tuple[int, ...] = (), enable_top_p_first: bool = False)[source]#

Bases: NodeResourceSpec

Parameters:

  • resource_key (str)

  • nodes (set[str])

  • vocab_size (int | None)

  • enable_repetion_penalty (bool)

  • enable_min_p (bool)

  • max_repetition_window (int)

  • min_tokens_stop_ids (tuple[int, ...])

  • enable_top_p_first (bool)

enable_min_p: bool = False#
enable_repetion_penalty: bool = True#
enable_top_p_first: bool = False#
max_repetition_window: int = 0#
min_tokens_stop_ids: tuple[int, ...] = ()#
property resource_class: type[Resource]#

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

vocab_size: int | None#
class mstar.engine.resources.SamplerStep(segments: tuple[mstar.engine.resources.step.Segment, ...] | None=None, apply_penalty: bool = True, apply_filters: bool = True, prefill_tracked_tokens: dict[str, torch.Tensor]=<factory>, kept_rids: frozenset[int] | None = None)[source]#

Bases: ResourceStep

Parameters:
apply_filters: bool = True#
apply_penalty: bool = True#
kept_rids: frozenset[int] | None = None#
prefill_tracked_tokens: dict[str, Tensor]#
class mstar.engine.resources.SamplingReqConfig(temperature: float = 0.6, top_k: int = 0, top_p: float = 1, ignore_eos: bool = False, repetition_penalty: float = 1, min_p: float = 0.0, repetition_window: int = 0, min_tokens: int = 0, top_p_first: bool = False, top_p_min_keep: int = 1, penalize_prompt: bool = True, _seed: int = 0)[source]#

Bases: ResourceReqConfig

Parameters:
apply_conductor_config(seed=0, resource_key=None, **kwargs)[source]#
Parameters:

  • seed (int)

  • resource_key (str | None)

ignore_eos: bool = False#
min_p: float = 0.0#
min_tokens: int = 0#
penalize_prompt: bool = True#
repetition_penalty: float = 1#
repetition_window: int = 0#
property seed#
temperature: float = 0.6#
top_k: int = 0#
top_p: float = 1#
top_p_first: bool = False#
top_p_min_keep: int = 1#
validate(spec=None)[source]#

Raise ValueError on a value outside the contract, or one spec’s node cannot honour. None means the sampler’s default.

Parameters:

spec (SamplerSpec | None)

Return type:

None

class mstar.engine.resources.Segment(request_id, label, span)[source]#

Bases: NamedTuple

One step’s addition to a request’s cache stream.

A request contributes one segment per label active for it in a step; the batch’s ordered segment list defines the layout of per-token arrays. span may be 0: a zero-span segment reads its stream without extending it (admission reserves nothing, commit is a no-op).

A NamedTuple, not a frozen dataclass: one is built per request per step, and the frozen dataclass’s object.__setattr__-per-field __init__ is the expensive way to do that.

Parameters:
label: str#

Alias for field number 1

request_id: str#

Alias for field number 0

span: int#

Alias for field number 2

class mstar.engine.resources.SlotLease(slot, bucket)[source]#

Bases: object

attention admit gives to inform which slot to plan and replay

has no clean channel to plan/commit/release. see O.

Parameters:
bucket: BucketKey | None#
slot: int#
class mstar.engine.resources.StepContext(request_ids: collections.abc.Sequence[int], graph_walk: str, slot: int, capture: bool, is_preplan: bool = False, plan_results: dict[str, typing.Any] = <factory>, slot_lease: mstar.engine.resources.step.SlotLease | None = None, piecewise_leases: 'Mapping[str, SlotLease]' = <factory>, _padded_request_ids: collections.abc.Sequence[int] | None = None, _real_rids: frozenset | None = None, request_walks: collections.abc.Mapping[int, str] = <factory>, force_prefill: bool = False)[source]#

Bases: object

Parameters:
capture: bool#
force_prefill: bool = False#
graph_walk: str#
is_padding_row(rid)[source]#

Whether rid is a replay’s padding row rather than a real request.

It carries the template span so shapes match the capture, but a resource should reserve, address and commit nothing on its behalf.

Parameters:

rid (str)

Return type:

bool

is_preplan: bool = False#
is_real_row(rid)[source]#

Whether rid is one of this step’s real requests.

Parameters:

rid (str)

Return type:

bool

property padded_request_ids: Sequence[int]#
piecewise_leases: Mapping[str, SlotLease]#
plan_results: dict[str, Any]#
request_ids: Sequence[int]#
request_walks: Mapping[int, str]#
set_padded_rids(padded_rids)[source]#
Parameters:

padded_rids (Sequence[int] | None)

set_piecewise_leases(leases)[source]#
Parameters:

leases (Mapping[str, SlotLease])

slot: int#
slot_lease: SlotLease | None = None#
walk_of(rid)[source]#
Parameters:

rid (int)

Return type:

str

class mstar.engine.resources.StepRunner(resources, node_resources=None, enable_nvtx=False)[source]#

Bases: object

drives resources through per-step cycle

Parameters:

  • resources (Mapping[str, Resource])

  • node_resources (Mapping[str, Collection[str]] | None)

  • enable_nvtx (bool)

admit(step)[source]#

Reserve capacity for step. A refused admit gives back everything the step reserved.

Parameters:

step (SubmoduleStep)

Return type:

FullAdmitOutcome

admit_retrieve(rid, node_name, graph_walk, published=None)[source]#

bring published state in

gives ready=False when still has inflights; shortcircuit on failure

Swept over node_name’s own resources. Asking a node about resources it does not own is not just wasted work: KVManager.get_labels answers for an unrecognised node with the default [“main”], so the node ends up gated on — and able to allocate against — another node’s cache.

Parameters:
Return type:

FullAdmitOutcome

adopt_cached_prefix(rid, node_name, graph_walk, leases, caught_up)[source]#
Parameters:
Return type:

bool

apply_cached_prefix(rid, node_name, graph_walk, inputs, prefix)[source]#

Give every resource the agreed prefix, with the untrimmed inputs.

Parameters:
Return type:

None

build_accelerator_graph_buffers(slots, max_bs, max_seq_len, node_name=None)[source]#

resources preallocate the static buffers captured replays will read

Swept over node_name’s own resources: another node’s slot count and batch size say nothing about a resource this node’s graphs never touch.

Parameters:
Return type:

None

clear_preplan()[source]#

Drop the staged pre-plan on every resource, and the record of it.

Return type:

None

commit(step)[source]#

record step consumption

Parameters:

step (SubmoduleStep)

Return type:

None

complete_cached_walk(rid, node_name, graph_walk)[source]#

Tell the node’s resources that the cache served this walk whole.

Parameters:
Return type:

None

extend_prefix_chains(rid, node_name, graph_walk, outputs, keys=None)[source]#

Offer this step’s sampled tokens to the node’s own resources.

Parameters:
Return type:

None

extend_prefix_chains_batch(request_ids, node_name, walk_of, host_rows, outputs, keys)[source]#

extend_prefix_chains for a step: a resource with a batch form takes the host rows; the others are offered each request’s host copy.

walk_of rather than one walk for the step, so a combined walk comes through here too instead of falling back to a call per row; a step of one walk passes that walk itself.

Parameters:
Return type:

None

ingest_capture_row(rid)[source]#

open a capture padding row on all resources

Parameters:

rid (str)

Return type:

None

ingest_request(rid, overrides=None)[source]#

open state on all resources; overrides are per-resource. All or nothing: a resource that raises leaves the request on none of them.

Parameters:
Return type:

None

lease_cached_prefix(rid, node_name, graph_walk)[source]#
Parameters:
Return type:

dict[str, list[int]]

property order: tuple[str, ...]#
page_trails(node_name)[source]#
Parameters:

node_name (str)

Return type:

dict[str, tuple[int, int]]

plan(step)[source]#

plan in dependency order

place plan in step.ctx.plan_results before next plan runs again could possibly move that into plan itself

Parameters:

step (SubmoduleStep)

Return type:

dict[str, Any]

pre_admit(step)[source]#

admit over the pre-planning subset, a step ahead

the later full admit covers the rest; these resources see their own state as already reserved and no-op

Unwound on refusal exactly as admit is.

Parameters:

step (SubmoduleStep)

Return type:

FullAdmitOutcome

pre_plan(step)[source]#

plan the pre-planning subset, a step ahead

ctx.is_preplan sends each one’s output to its pending slot rather than the live one; the later full plan promotes it. Results are left in ctx.plan_results for the dependents in this same subset.

Parameters:

step (SubmoduleStep)

Return type:

dict[str, Any]

prefix_chain_keys(node_name)[source]#

The node’s resources with a prefix cache open right now. Empty in the common case, which lets the caller skip a per-request sweep that was six no-op calls per request per step.

Parameters:

node_name (str | None)

Return type:

list[str]

publish(request_ids, node_name=None, graph_walk=None, request_walks=None)[source]#

durable state outward publish

non-publish resources noop; pairs with admit_retrieve

Scoped to the node whose step just ran; no other node’s state moved, and merge_publish_info keeps the entry it already published.

Parameters:
Return type:

dict[str, dict[str, PublishedInfo]]

publish_after_stop(request_ids, node_name=None, graph_walk=None, request_walks=None)[source]#

Publish resources configured for the end of a dynamic loop.

Parameters:
Return type:

dict[str, dict[str, PublishedInfo]]

publish_from_snapshots(snapshots, request_ids, node_name=None, graph_walk=None, request_walks=None)[source]#

The published info of request_ids from their snapshots, each under its own walk (see publish_snapshot).

Parameters:
Return type:

dict[str, dict[str, PublishedInfo]]

publish_snapshot(request_ids, node_name=None, graph_walk=None, request_walks=None)[source]#

publish, deferred: per rid the snapshot each publisher took.

Finished by publish_from_snapshots for the rids that need it. A rid with nothing to publish has no entry. Under a combined walk each rid is snapshotted under its own walk (request_walks), as publish does: a publisher may key what it publishes on the walk.

Parameters:
Return type:

dict[str, dict[str, Any]]

release_cached_prefix(rid, node_name, graph_walk)[source]#
Parameters:
Return type:

None

remove_request(rid)[source]#
Parameters:

rid (str)

Return type:

None

resolve_cached_prefix(rid, node_name, graph_walk, inputs)[source]#

The longest prefix of this walk every resource on node_name can serve.

The smallest answer wins, because a step is one span: a resource that holds less than another has to be caught up by the tokens that run. A resource with no opinion is not an answer of zero.

Parameters:
Return type:

CachedPrefix | None

property resources: Mapping[str, Resource]#
rollback_admit(step)[source]#

Give back everything an admitted step reserved, for a step that will not run: the unwind a refused admit does, over every resource the step declares.

Parameters:

step (SubmoduleStep)

Return type:

None

property staged: bool#

A pre-plan is staged and not yet promoted.

class mstar.engine.resources.SubmoduleStep(steps: dict[str, mstar.engine.resources.step.ResourceStep], segments: list[mstar.engine.resources.step.Segment] | None = None, cg_key_info: Any | None = None, _ctx: mstar.engine.resources.step.StepContext = None)[source]#

Bases: object

Parameters:
cg_key_info: Any | None = None#
property ctx#
get(key)[source]#
Parameters:

key (str)

Return type:

ResourceStep | None

keys()[source]#
Return type:

KeysView[str]

segments: list[Segment] | None = None#
set_ctx(ctx)[source]#
steps: dict[str, ResourceStep]#
mstar.engine.resources.apply_yaml_overrides(specs, model_config)[source]#

Apply a deployment’s resources: block to the specs it names.

resources: {<resource_key>: {...}}, one block per resource, so a model with two pools of the same kind (whisper’s decoder cache and its encoder context) can have each tuned on its own. An unknown key is an error: it would otherwise be a silently ineffective setting.

Parameters:
Return type:

None

mstar.engine.resources.cross_label(q_label, kv_label)[source]#

The plan a cross-attention pair runs under: q_label’s spans attending kv_label’s.

Parameters:

  • q_label (str)

  • kv_label (str)

Return type:

str

mstar.engine.resources.resolve_spec_dependencies(specs)[source]#

Index specs by resource key, checking uniqueness and depends_on.

Parameters:

specs (Sequence[NodeResourceSpec])

Return type:

dict[str, NodeResourceSpec]

mstar.engine.resources.topo_sort(resources)[source]#

Keys in dependency order; built resources at load, their specs when the conductor validates a deployment.

Parameters:

resources (Mapping[str, Resource | NodeResourceSpec])

Return type:

tuple[str, …]

Modules

attn

base

kv

position

recurrent

rms_norm

Backend-neutral RMSNorm operator.

runner

sequencing for resource call cycle

sampler

spec

Resource declarations models hand to the engine.

step

Contents