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#
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)[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.

Parameters:

message (str)

class mstar.engine.resources.AllocationFailed(message: str, pages_short: int, label: str, request_id: str)[source]#

Bases: AdmitFailedReason

Parameters:

  • message (str)

  • pages_short (int)

  • label (str)

  • request_id (str)

label: str#
pages_short: int#
request_id: str#
class mstar.engine.resources.AttentionConfig(kv_cache: str, backend: mstar.engine.resources.attn.config.AttnBackend = <AttnBackend.FLASHINFER: 'flashinfer'>, flashinfer_backend: str = 'auto')[source]#

Bases: object

Parameters:
backend: AttnBackend = 'flashinfer'#
flashinfer_backend: str = 'auto'#
kv_cache: str#
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)[source]#

Bases: ResourceStep

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

Bases: Enum

DENSE = 'dense'#
FLASHINFER = 'flashinfer'#
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: 'CudaGraphConfig | PiecewiseCudaGraphConfig', config_idx: 'int | None' = None)[source]#

Bases: object

Parameters:
property bs#
bucket: BucketKey#
config: CudaGraphConfig | PiecewiseCudaGraphConfig#
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.FLASHINFER, 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 = 'flashinfer'#
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

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)[source]#

Bases: ResourceReqConfig

Parameters:
apply_conductor_config(prefix_keys=None, prefix_tail=None, prefix_decode=None, prefix_cache=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]

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]]#
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#
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], ...] = (), 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#
post_forks: tuple[tuple[str, str], ...] = ()#
pre_forks: tuple[tuple[str, str], ...] = ()#
retention: Mapping[tuple[str, str], RetentionPolicy]#
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)[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)

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)[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)

Return type:

None

cpu_offload_pages: int = 0#
layout: KVLayout = 'NHD'#
max_num_pages: int = 2048#
max_seq_len: int#
page_size: int = 128#
prefix_cache: bool = True#
prefix_cache_salt: str = ''#
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, 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_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, flashinfer_backend='auto', dtype=None)[source]#

Bases: object

Varlen self-attention over segments packed into one forward, with no KV cache: the whole layout is this step’s, and nothing carries to the next.

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

Parameters:

  • num_qo_heads (int)

  • num_kv_heads (int)

  • head_dim (int)

  • sm_scale (float | None)

  • max_segments_per_request (int)

  • max_tokens_per_request (int | None)

  • flashinfer_backend (str)

  • dtype (dtype | None)

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: str, nodes: set[str], config: mstar.engine.resources.attn.ragged.config.RaggedAttentionConfig)[source]#

Bases: NodeResourceSpec

Parameters:
apply_yaml_overrides(flashinfer_backend=None, max_segments_per_request=None, max_tokens_per_request=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.

The two ceilings are here rather than on the model because they size CUDA-graph buckets, which is a deployment’s memory/coverage trade.

Parameters:

  • flashinfer_backend (str | None)

  • max_segments_per_request (int | None)

  • max_tokens_per_request (int | None)

config: RaggedAttentionConfig#
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.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

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

Take on the length every resource agreed to skip.

The inputs are untrimmed, so a resource that needs the tokens being skipped can still read them.

Parameters:

  • rid (str)

  • node_name (str)

  • graph_walk (str)

  • inputs (Any)

  • matched_len (int)

Return type:

None

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

Resource

build_cuda_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_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

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.

walks names, per label, the walk that writes the keyed span and the one that decodes after it.

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_request(rid, overrides)[source]#
Parameters:
is_offloaded(rid)[source]#
Parameters:

rid (str)

Return type:

bool

offload(rid)[source]#

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

Parameters:

rid (str)

Return type:

int

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

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

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)[source]#

How many of this request’s leading tokens this resource already holds.

None leaves the length to other resources; 0 means nothing can be skipped.

Parameters:
Return type:

int | 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 hook below.

apply_conductor_config(**kwargs)[source]#
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)[source]#

Bases: NodeResourceSpec

Parameters:

  • resource_key (str)

  • nodes (set[str])

  • vocab_size (int | None)

  • enable_repetion_penalty (bool)

  • enable_min_p (bool)

enable_min_p: bool = False#
enable_repetion_penalty: bool = True#
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, prefill_tracked_tokens: dict[str, torch.Tensor]=<factory>)[source]#

Bases: ResourceStep

Parameters:
apply_penalty: bool = True#
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, _seed: int = 0)[source]#

Bases: ResourceReqConfig

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

seed (int)

ignore_eos: bool = False#
min_p: float = 0.0#
repetition_penalty: float = 1#
property seed#
temperature: float = 0.6#
top_k: int = 0#
top_p: float = 1#
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)[source]#

Bases: object

Parameters:
capture: bool#
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#
property padded_request_ids: Sequence[int]#
piecewise_leases: Mapping[str, SlotLease]#
plan_results: dict[str, Any]#
request_ids: Sequence[int]#
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#
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

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

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

Parameters:

  • rid (str)

  • node_name (str)

  • graph_walk (str)

  • matched_len (int)

Return type:

None

build_cuda_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

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

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

Parameters:
Return type:

None

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

open state on all resources; overrides are per-resource

Parameters:
Return type:

None

property order: tuple[str, ...]#
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]

publish(request_ids, node_name=None, graph_walk=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:

  • request_ids (list[str])

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

dict[str, dict[str, PublishedInfo]]

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

Publish resources configured for the end of a dynamic loop.

Parameters:

  • request_ids (list[str])

  • node_name (str | None)

  • graph_walk (str | None)

Return type:

dict[str, dict[str, PublishedInfo]]

remove_request(rid)[source]#
Parameters:

rid (str)

Return type:

None

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

The longest prefix 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:

int

property resources: Mapping[str, Resource]#
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.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]#
Parameters:

resources (Mapping[str, Resource])

Return type:

tuple[str, …]

Modules

attn

base

convenience

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