mstar.engine.resources.recurrent

Contents

mstar.engine.resources.recurrent#

class mstar.engine.resources.recurrent.DeltaNetGeometry(num_k_heads, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size)[source]#

Bases: RecurrentGeometry

Head geometry of the delta-net family: gated delta rule (Qwen3.5, Qwen3-Next) and Kimi delta attention (Kimi Linear, GLM-5.3).

Both carry a K-last [HV, V, K] state matrix (FlashInfer’s pool layout) and a conv window holding every tap but the current token’s. to_blocks and from_blocks are inverses, so a backend reads its geometry off the pool.

Parameters:

  • num_k_heads (int)

  • num_v_heads (int)

  • head_k_dim (int)

  • head_v_dim (int)

  • conv_kernel_size (int)

property conv_dim: int#

The depthwise conv runs over [q | k | v] concatenated.

conv_kernel_size: int#
classmethod from_blocks(blocks)[source]#

Recover head counts from (possibly sharded) block shapes.

Raises if the shapes are not a delta-net’s.

Parameters:

blocks (dict[str, RecurrentBlockConfig])

Return type:

DeltaNetGeometry

head_k_dim: int#
head_v_dim: int#
num_k_heads: int#
num_v_heads: int#
to_blocks(state_dtype=torch.float32, conv_dtype=torch.bfloat16)[source]#

Pool blocks for this geometry; head counts are pre-sharding.

Parameters:
Return type:

dict[str, RecurrentBlockConfig]

class mstar.engine.resources.recurrent.Mamba2Geometry(num_heads, head_dim, state_size, n_groups, conv_kernel_size)[source]#

Bases: RecurrentGeometry

Head geometry of a Mamba-2 (SSD) layer: Nemotron-H, NVIDIA’s Nemotron Nano/VoiceChat backbones, Mamba-Codestral.

Two blocks, like the delta-net family: an [H, P, N] SSM state per head (fp32 by default: it is an exponentially decayed running sum over the whole session, and the reference keeps it in fp32) and the conv window over the [x | B | C] projection, every tap but the current token’s.

n_groups is the number of B/C groups shared by H heads. It is a shape fact (conv_dim = H * P + 2 * n_groups * N), so it is recoverable from the blocks and from_blocks inverts to_blocks.

Parameters:

  • num_heads (int)

  • head_dim (int)

  • state_size (int)

  • n_groups (int)

  • conv_kernel_size (int)

property conv_dim: int#

The depthwise conv runs over [x | B | C] concatenated.

conv_kernel_size: int#
property d_inner: int#
classmethod from_blocks(blocks)[source]#

Recover the geometry from block shapes (sharded ones too, since every axis involved shards). Raises on shapes that are not a Mamba-2’s, which is the check that a pool and the resource planning against it were built for the same model.

Parameters:

blocks (dict[str, RecurrentBlockConfig])

Return type:

Mamba2Geometry

head_dim: int#
n_groups: int#
num_heads: int#
state_size: int#
to_blocks(state_dtype=torch.float32, conv_dtype=torch.bfloat16)[source]#

Pool blocks for this geometry; head counts pre-sharding (heads and the conv channels both shard on their leading axis).

Parameters:
Return type:

dict[str, RecurrentBlockConfig]

class mstar.engine.resources.recurrent.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.RecurrentBlockConfig(shape, dtype, shard_dims=())[source]#

Bases: object

One per-slot, per-layer tensor block.

shape is opaque to the pool. shard_dims names the axes divided across ranks.

Parameters:
dtype: dtype#
property nbytes: int#
property numel: int#
shape: tuple[int, ...]#
shard(num_shards)[source]#
Parameters:

num_shards (int)

Return type:

None

shard_dims: tuple[int, ...] = ()#
class mstar.engine.resources.recurrent.RecurrentGeometry[source]#

Bases: ABC

abstractmethod classmethod from_blocks(blocks)[source]#
Parameters:

blocks (dict[str, RecurrentBlockConfig])

Return type:

RecurrentGeometry

abstractmethod to_blocks(*args, **kwargs)[source]#
Return type:

dict[str, RecurrentBlockConfig]

class mstar.engine.resources.recurrent.RecurrentStateConfig(num_layers: int, blocks: dict[str, mstar.engine.resources.recurrent.config.RecurrentBlockConfig]=<factory>, max_slots: int = 256, disable_sink_slot: bool = False)[source]#

Bases: object

Parameters:
blocks: dict[str, RecurrentBlockConfig]#
disable_sink_slot: bool = False#
max_slots: int = 256#
num_layers: int#
shard(num_shards)[source]#

Narrow every block’s sharded axes. Idempotent, so the pool and its resource can both shard the shared config.

Parameters:

num_shards (int)

Return type:

None

property slot_bytes: int#
property total_bytes: int#
property usable_slots: int#

Slots requests can hold; the sink is not one of them.

class mstar.engine.resources.recurrent.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.RecurrentStateSpec(resource_key: str, nodes: set[str], config: mstar.engine.resources.recurrent.config.RecurrentStateConfig)[source]#

Bases: NodeResourceSpec

Parameters:
apply_yaml_overrides(max_slots=None, state_dtype=None)[source]#

How many slots this deployment gets, and the state’s dtype.

Block shapes are the model’s and not tunable. The dtype trades accumulated precision against bandwidth and decides which kernels the backend can reach.

Parameters:

  • max_slots (int | None)

  • state_dtype (str | None)

config: RecurrentStateConfig#
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.recurrent.RecurrentStep(segments=None, pre_forks=(), post_forks=())[source]#

Bases: ResourceStep

One step’s work against the pool.

No commit flag, unlike KVStep: a backend writes the pool in place, so a non-committing reader needs two labels, reading one and writing another.

Forks mirror KVStep’s: (from_label, to_label) pairs, reserved at admit and copied at plan (pre) or commit (post).

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

Modules

config

What a model declares about a pool of recurrent state.

pool

A pool of fixed-size recurrent state slots.
Contents