mstar.engine.resources.base

Contents

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

CachedPrefix(label, tokens, position)

How much of a walk's input a resource already holds for label.

EngineResourceInfo(device[, ...])

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

KeyedWalks(writes[, growth])

A keyed label's walks, as a resource's prefix cache takes them: the walks whose writes are keyed, and the growth walks whose sampled ids extend the keys.

PublishedInfo()

Resource()

class mstar.engine.resources.base.AttentionResource[source]#

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()[source]#
Return type:

None

set_default_label(label)[source]#
Parameters:

label (str)

Return type:

None

set_default_layer_idx(layer_idx)[source]#
Parameters:

layer_idx (int)

Return type:

None

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

Bases: object

Parameters:
bucket: BucketKey#
label: str#
slot: int#
class mstar.engine.resources.base.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.base.CachedPrefix(label, tokens, position)[source]#

Bases: NamedTuple

How much of a walk’s input a resource already holds for label.

tokens counts the walk’s leading inputs held, and position is the stream’s position counter once they are skipped.

Parameters:
label: str#

Alias for field number 0

position: int#

Alias for field number 2

tokens: int#

Alias for field number 1

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>, max_rows_in_combined=<factory>)[source]#

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:
dependencies: Mapping[str, NodeResourceSpec]#
dependency(key)[source]#
Parameters:

key (str)

Return type:

NodeResourceSpec

device: torch.device#
joint_comm_group: JointGroups | None = None#
kv_dtype: torch.dtype = torch.bfloat16#
max_rows_in_combined: Mapping[str, int | None]#
needs_remote_transfer: bool = True#
transfer_engine_info: TransferEngineInfo | None = None#
class mstar.engine.resources.base.KeyedWalks(writes, growth=())[source]#

Bases: NamedTuple

A keyed label’s walks, as a resource’s prefix cache takes them: the walks whose writes are keyed, and the growth walks whose sampled ids extend the keys.

Parameters:
growth: tuple[str, ...]#

Alias for field number 1

writes: tuple[str, ...]#

Alias for field number 0

class mstar.engine.resources.base.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.base.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#
mstar.engine.resources.base.build_resource(spec, info)[source]#
Parameters:
Return type:

Resource

Contents