mstar.engine.resources

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)
ready (bool)
reason (AdmitFailedReason | None)

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:

kv_cache (str)
backend (AttnBackend)
flashinfer_backend (str)

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:

resource_key (str)
nodes (set[str])
config (AttentionConfig)

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:

backend (str | AttnBackend | None)
flashinfer_backend (str | None)

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:

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

causal: bool = True#

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

Bases: Enum

DENSE = 'dense'#

FLASHINFER = 'flashinfer'#

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:

bucket (BucketKey)
slot (int)
config (CudaGraphConfig | PiecewiseCudaGraphConfig)
config_idx (int | None)

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:

kv_cache (str)
query_kv_cache (str | None)
context_label (str)
backend (AttnBackend)
flashinfer_backend (str)

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:

resource_key (str)
nodes (set[str])
config (CrossAttentionConfig)

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

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

Parameters:

backend (str | AttnBackend | None)
flashinfer_backend (str | None)

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:

outcome (AdmitOutcome)
failed_resource (str | None)

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: int, num_kv_heads: int, head_dim: int, max_seq_len: int, max_num_pages: int = 2048, page_size: int = 128, num_qo_heads: int = None, layout: mstar.engine.resources.kv.config.KVLayout = <KVLayout.NHD: 'NHD'>, cpu_offload_pages: int = 0)[source]#

Bases: object

Parameters:

num_layers (int)
num_kv_heads (int)
head_dim (int)
max_seq_len (int)
max_num_pages (int)
page_size (int)
num_qo_heads (int)
layout (KVLayout)
cpu_offload_pages (int)

cpu_offload_pages: int = 0#

head_dim: int#

layout: KVLayout = 'NHD'#

max_num_pages: int = 2048#

max_seq_len: int#

num_kv_heads: int#

num_layers: int#

num_qo_heads: int = None#

page_size: int = 128#

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

Bases: ResourceReqConfig

Parameters:

needed_labels (list[str] | None)
needed_labels_per_node (dict[str, list[str]])
needed_labels_per_node_walk (dict[tuple[str, str], list[str]])

get_labels(node, walk)[source]#

Parameters:

node (str)
walk (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]]#

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

Bases: NodeResourceSpec

Parameters:

resource_key (str)
nodes (set[str])
config (KVConfig)

apply_yaml_overrides(max_num_pages=None, page_size=None, max_seq_len=None, cpu_offload_pages=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)

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], ...] = ())[source]#

Bases: ResourceStep

Parameters:

segments (tuple[Segment, ...] | None)
commit (bool)
combined_labels (dict[tuple[str, ...], str])
pre_forks (tuple[tuple[str, str], ...])
post_forks (tuple[tuple[str, str], ...])

combined_labels: dict[tuple[str, ...], str]#

commit: bool = True#

post_forks: tuple[tuple[str, str], ...] = ()#

pre_forks: tuple[tuple[str, str], ...] = ()#

class mstar.engine.resources.NodeResourceSpec(resource_key: str, nodes: set[str])[source]#

Bases: ABC

Parameters:

resource_key (str)
nodes (set[str])

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

kv_cache (str)
backend (PosBackend)
scheme (PosScheme)
block_step (int)
rotary_dim (int | None)
interleave (bool)
rope_scale (float)
rope_theta (float)
rope_dtype (dtype | None)
low_freq_factor (float | None)
high_freq_factor (float | None)
old_context_len (int | None)

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:

resource_key (str)
nodes (set[str])
config (PositionConfig)

config: PositionConfig#

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:

segments (tuple[Segment, ...] | None)
pos_ids (dict[str, Tensor] | Tensor | None)
advance (tuple[int, ...] | None)

advance: tuple[int, ...] | None = None#

pos_ids: dict[str, Tensor] | Tensor | None = None#

class mstar.engine.resources.PublishedInfo[source]#

Bases: ABC

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

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:

resource_key (str)
nodes (set[str])
config (RaggedAttentionConfig)

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:

message (str)
label (str)
request_id (str)

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:

step (ResourceStep)
ctx (StepContext)

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:

rid (str)
node_name (str)
graph_walk (str)
published (PublishedInfo | None)

Return type:

AdmitOutcome

abstractmethod classmethod build(spec, info)[source]#

Parameters:

spec (NodeResourceSpec)
info (EngineResourceInfo)

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:

slots (list[CGSlotSpec])
max_bs (int)
max_seq_len (int)

Return type:

None

cleanup()[source]#

clear_preplan()[source]#

commit(step, ctx)[source]#

record step consumption

Parameters:

step (ResourceStep)
ctx (StepContext)

Return type:

None

depends_on()[source]#

Return type:: set[str]

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:

rid (str)
overrides (ResourceReqConfig | None)

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:

step (ResourceStep)
ctx (StepContext)

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.

publish(request_id)[source]#

Parameters:: request_id (str)
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:

rid (str)
free (bool)

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.SamplerSpec(resource_key: str, nodes: set[str], vocab_size: int | None, enable_repetion_penalty: bool = True)[source]#

Bases: NodeResourceSpec

Parameters:

resource_key (str)
nodes (set[str])
vocab_size (int | None)
enable_repetion_penalty (bool)

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:

segments (tuple[Segment, ...] | None)
apply_penalty (bool)
prefill_tracked_tokens (dict[str, Tensor])

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

Bases: ResourceReqConfig

Parameters:

temperature (float)
top_k (int)
top_p (float)
ignore_eos (bool)
repetition_penalty (float)
_seed (int)

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

Parameters:: seed (int)

ignore_eos: bool = False#

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:

request_id (str)
label (str)
span (int)

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:

slot (int)
bucket (BucketKey | None)

bucket: BucketKey | None#

slot: int#

class mstar.engine.resources.StepContext(request_ids: collections.abc.Sequence[str], 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[str] | None = None)[source]#

Bases: object

Parameters:

request_ids (Sequence[str])
graph_walk (str)
slot (int)
capture (bool)
is_preplan (bool)
plan_results (dict[str, Any])
slot_lease (SlotLease | None)
piecewise_leases (Mapping[str, SlotLease])
_padded_request_ids (Sequence[str] | None)

capture: bool#

graph_walk: str#

is_preplan: bool = False#

property padded_request_ids: Sequence[str]#

piecewise_leases: Mapping[str, SlotLease]#

plan_results: dict[str, Any]#

request_ids: Sequence[str]#

set_padded_rids(padded_rids)[source]#

Parameters:: padded_rids (Sequence[str] | 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

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:

rid (str)
node_name (str)
graph_walk (str)
published (Mapping[str, PublishedInfo] | None)

Return type:

FullAdmitOutcome

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:

slots (list[CGSlotSpec])
max_bs (int)
max_seq_len (int)
node_name (str | None)

Return type:

None

commit(step)[source]#

record step consumption

Parameters:: step (SubmoduleStep)
Return type:: None

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

open state on all resources; overrides are per-resource

Parameters:

rid (str)
overrides (Mapping[str, ResourceReqConfig] | None)

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

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

Return type:

dict[str, dict[str, PublishedInfo]]

remove_request(rid)[source]#

Parameters:: rid (str)
Return type:: None

property resources: Mapping[str, Resource]#

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:

steps (dict[str, ResourceStep])
segments (list[Segment] | None)
cg_key_info (Any | None)
_ctx (StepContext)

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:

specs (Sequence[NodeResourceSpec])
model_config (Mapping[str, Any])

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`
`runner`	sequencing for resource call cycle
`sampler`
`spec`	Resource declarations models hand to the engine.
`step`

mstar.engine.resources

Contents

mstar.engine.resources#