mstar.engine.resources.recurrent.config

mstar.engine.resources.recurrent.config#

What a model declares about a pool of recurrent state.

Kept free of the manager so declaring a step does not import a backend. The pool is ignorant of what the state means: a slot is a fixed number of bytes per layer, so capacity is a slot count and a fork is a fixed-size copy.

Classes

DeltaNetGeometry(num_k_heads, num_v_heads, ...)

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

Mamba2Geometry(num_heads, head_dim, ...)

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

RecurrentBlockConfig(shape, dtype[, shard_dims])

One per-slot, per-layer tensor block.

RecurrentGeometry()

RecurrentStateConfig(num_layers, blocks, ...)

RecurrentStateSpec(resource_key, nodes, config)

RecurrentStep([segments, pre_forks, post_forks])

One step's work against the pool.
class mstar.engine.resources.recurrent.config.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.config.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.config.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.config.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.config.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.config.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.config.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], ...] = ()#