mstar.engine.resources.recurrent.pool

mstar.engine.resources.recurrent.pool#

A pool of fixed-size recurrent state slots.

Storage only: the pool hands out slots, keeps the per-layer tensors they index, and says which slot each row of a step addresses. As kv/ is to attn/, a linear-attention resource is the wrapper planned against it.

Classes

RecurrentAddressing(slot_indices, has_state, ...)

Where one label's rows live this step: the pool's whole plan output.

RecurrentStatePool(device, config)

SlotState(index[, has_state, generation])

class mstar.engine.resources.recurrent.pool.RecurrentAddressing(slot_indices, has_state, num_rows, host_slot_indices=None, host_has_state=None)[source]#

Bases: object

Where one label’s rows live this step: the pool’s whole plan output.

Token layout (cu_seqlens, chunk metadata) is the backend’s to build.

Parameters:
has_state: Tensor#
host_has_state: Sequence[bool] | None = None#
host_slot_indices: Sequence[int] | None = None#
num_rows: int#
slot_indices: Tensor#
class mstar.engine.resources.recurrent.pool.RecurrentStatePool(device, config)[source]#

Bases: Resource

Parameters:
addressing(label='main')[source]#

This step’s addressing under label, for a backend to run against.

Parameters:

label (str)

Return type:

RecurrentAddressing

admit(step, ctx)[source]#

Reserve a slot per addressed (rid, label), plus fork targets.

Capture rows and a replay’s padding rows reserve nothing and address the sink instead; otherwise the runner’s dummy names would drain a pool sized for the real concurrency.

Parameters:
Return type:

AdmitOutcome

block(name, layer_idx)[source]#

One layer’s contiguous slot-major view of a block: [max_slots, *shape].

Parameters:
Return type:

Tensor

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

RecurrentStatePool

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

cleanup()[source]#
clear_preplan()[source]#

Drop the staged plan; staging touched no live state, so nothing to rewind.

commit(step, ctx)[source]#

record step consumption

Parameters:
Return type:

None

ingest_request(rid, overrides=None)[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:
property num_free_slots: int#
property pad_index: int#

What an unaddressed row points at.

plan(step, ctx)[source]#

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

Parameters:
remove_request(rid)[source]#
Parameters:

rid (str)

reset_request(rid, free=False)[source]#

Drop a request’s slots. free hands them back to the pool.

Both paths zero what they release, or a backend could resume on another request’s state.

Parameters:
rollback_admit(step, ctx)[source]#

Free the slots this step’s admit took fresh; slots a request already held stay with it. Without this, two forking requests can each keep the slot the other’s fork needs, and the pool has no eviction to break that.

Parameters:
Return type:

None

property supports_preplan#
class mstar.engine.resources.recurrent.pool.SlotState(index: 'int', has_state: 'bool' = False, generation: 'int' = 0)[source]#

Bases: object

Parameters:
generation: int = 0#
has_state: bool = False#
index: int#