mstar.engine.resources.base

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[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: '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.base.EngineResourceInfo(device, joint_comm_group=None, transfer_engine_info=None, kv_dtype=torch.bfloat16, dependencies=<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#
transfer_engine_info: TransferEngineInfo | None = None#
class mstar.engine.resources.base.PublishedInfo[source]#

Bases: ABC

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

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

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

record step consumption

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

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

Resource