mstar.engine.resources.base

Functions

build_resource(spec, info)

Classes

AttentionResource()	A resource a layer stack calls per layer, under one plan label.
CGSlotKey(bucket, slot, label)
CGSlotSpec(bucket, slot, config[, config_idx])
EngineResourceInfo(device[, ...])	What the engine has to offer a resource at build time.
PublishedInfo()
Resource()

class mstar.engine.resources.base.AttentionResource

Bases: Resource

A resource a layer stack calls per layer, under one plan label.

Adds the label / layer-index cursors: a caller running the whole stack sets them once instead of threading them through every call, and an explicit argument still supersedes. Class-level defaults so a subclass picks them up without touching its __init__; subclasses that read them clear them in plan, so a step that never binds cannot inherit the previous step’s.

The KV, attention and cross-attention resources; not the sampler or the position resource, which are called once per step rather than per layer.

property default_label: str

reset_default_cursors()

Return type:

None

set_default_label(label)

Parameters:

label (str)

Return type:

None

set_default_layer_idx(layer_idx)

Parameters:

layer_idx (int)

Return type:

None

class mstar.engine.resources.base.CGSlotKey(bucket: 'BucketKey', slot: 'int', label: 'str')

Bases: object

Parameters:

• bucket (BucketKey)

• slot (int)

• label (str)

bucket: BucketKey

label: str

slot: int

class mstar.engine.resources.base.CGSlotSpec(bucket: 'BucketKey', slot: 'int', config: 'CudaGraphConfig | PiecewiseCudaGraphConfig', config_idx: 'int | None' = None)

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.base.EngineResourceInfo(device, joint_comm_group=None, transfer_engine_info=None, kv_dtype=torch.bfloat16, needs_remote_transfer=True, dependencies=<factory>)

Bases: object

What the engine has to offer a resource at build time.

One struct rather than per-kind keyword arguments: a resource takes what it needs and ignores the rest, and a name that does not exist here is a TypeError rather than something silently swallowed by a **kwargs.

Parameters:

• device (torch.device)

• joint_comm_group (JointGroups | None)

• transfer_engine_info (TransferEngineInfo | None)

• kv_dtype (torch.dtype)

• needs_remote_transfer (bool)

• dependencies (Mapping[str, NodeResourceSpec])

dependencies: Mapping[str, NodeResourceSpec]

dependency(key)

Parameters:

key (str)

Return type:

NodeResourceSpec

device: torch.device

joint_comm_group: JointGroups | None = None

kv_dtype: torch.dtype = torch.bfloat16

needs_remote_transfer: bool = True

transfer_engine_info: TransferEngineInfo | None = None

class mstar.engine.resources.base.PublishedInfo

Bases: ABC

abstractmethod clone()

Return independent publication metadata for deferred sending.

Return type:

PublishedInfo

abstractmethod update(other)

Parameters:

other (PublishedInfo)

Return type:

None

class mstar.engine.resources.base.Resource

Bases: ABC

admit(step, ctx)

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)

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

apply_cached_prefix(rid, node_name, graph_walk, inputs, matched_len)

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)

Parameters:

• spec (NodeResourceSpec)

• info (EngineResourceInfo)

Return type:

Resource

build_cuda_graph_buffers(slots, max_bs, max_seq_len)

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

can_reload(rid)

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

clear_preplan()

commit(step, ctx)

record step consumption

Parameters:

• step (ResourceStep)

• ctx (StepContext)

Return type:

None

depends_on()

Return type:

set[str]

enable_prefix_cache(root, walks=None)

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:

• root (bytes)

• walks (dict[str, tuple[str, str | None]] | None)

Return type:

bool

extend_prefix_chain(rid, node_name, graph_walk, outputs)

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

Parameters:

• rid (str)

• node_name (str)

• graph_walk (str)

• outputs (Any)

Return type:

None

fingerprint()

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)

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)

Parameters:

• rid (str)

• overrides (ResourceReqConfig | None)

is_offloaded(rid)

Parameters:

rid (str)

Return type:

bool

offload(rid)

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

Parameters:

rid (str)

Return type:

int

plan(step, ctx)

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

Parameters:

• step (ResourceStep)

• ctx (StepContext)

Return type:

Any

post_warmup_validate()

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)

Parameters:

request_id (str)

Return type:

PublishedInfo | None

publish_after_stop(request_id, node_name, graph_walk)

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)

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)

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)

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

Parameters:

rid (str)

Return type:

bool

remove_request(rid)

Parameters:

rid (str)

reset_request(rid, free=False)

For clearing dummy RIDs during cuda graph capture

Parameters:

• rid (str)

• free (bool)

resolve_cached_prefix(rid, node_name, graph_walk)

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:

• rid (str)

• node_name (str)

• graph_walk (str)

Return type:

int | None

rollback_admit(step, ctx)

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:

• step (ResourceStep)

• ctx (StepContext)

Return type:

None

property supports_eviction

property supports_preplan

mstar.engine.resources.base.build_resource(spec, info)

Parameters:

• spec (NodeResourceSpec)

• info (EngineResourceInfo)

Return type:

Resource