mstar.engine.resources.convenience

mstar.engine.resources.convenience#

Classes

AttentionCallable(kv[, attn])

A convenience wrapper around kv and attn that wraps the KV write and attention call in one pure-tensor function, with helper methods for setting layer and label information.

LinearAttnCallable(pool[, attn])

A recurrent state pool plus the resource running kernels against it, mirroring AttentionCallable.

Mamba2Callable(pool[, attn])

LinearAttnCallable for a Mamba-2 layer: the conv, then the SSD recurrence, each handed this layer's block of the pool as a plain tensor.

RaggedAttentionCallable(attn, label)

AttentionCallable for a cacheless attention: (q, k, v) -> out over one declared span of a RaggedAttnManager.
class mstar.engine.resources.convenience.AttentionCallable(kv, attn=None)[source]#

Bases: object

A convenience wrapper around kv and attn that wraps the KV write and attention call in one pure-tensor function, with helper methods for setting layer and label information.

Must be used for ulysses_attention, which expects such a pure tensor callable. Recommended to use instance per transformer, not one per layer, as Dynamo specializes ulysses_attention on the identity of its run_attention argument, so a per-layer callable retraces that frame once per layer and blows the recompile limit.

That sharing is why the label is one cursor for the whole stack. No model varies its label per layer today; one that needs to should thread the label explicitly rather than use this.

For a cacheless attention, see RaggedAttentionCallable.

Parameters:
bind_step(label, attn=None)[source]#
Parameters:
Return type:

None

property label: str#

This step’s label. Read through to the resource, not stored here, so one instance can drive a whole stack of per-layer callables — and so a layer that no longer takes a label as an argument can still reach it (the position resource carries no cursor, so apply_qk is passed this).

set_layer_idx(layer_idx)[source]#
Parameters:

layer_idx (int)

Return type:

None

class mstar.engine.resources.convenience.LinearAttnCallable(pool, attn=None)[source]#

Bases: object

A recurrent state pool plus the resource running kernels against it, mirroring AttentionCallable.

The layer calls conv then the instance; this hands the pool’s per-layer blocks over as plain tensors, so neither the layer nor the manager holds the pool.

Parameters:
bind_step(label, attn=None)[source]#
Parameters:

  • label (str)

  • attn (LinearAttnManager | None)

Return type:

None

conv(x, weight, bias=None, activation='silu')[source]#
Parameters:
Return type:

Tensor

property label: str#
set_layer_idx(layer_idx)[source]#

The layer’s index among the recurrent layers, not the stack’s.

Parameters:

layer_idx (int)

Return type:

None

class mstar.engine.resources.convenience.Mamba2Callable(pool, attn=None)[source]#

Bases: object

LinearAttnCallable for a Mamba-2 layer: the conv, then the SSD recurrence, each handed this layer’s block of the pool as a plain tensor.

Same cursor protocol (bind_step once per stack, set_layer_idx per recurrent layer). The layer index counts recurrent layers only; a hybrid stack interleaves them with attention layers, and the pool is sized by its own count.

Parameters:
bind_step(label, attn=None)[source]#
Parameters:

  • label (str)

  • attn (LinearAttnManager | None)

Return type:

None

conv(x, weight, bias=None, activation='silu')[source]#
Parameters:
Return type:

Tensor

property label: str#
set_layer_idx(layer_idx)[source]#
Parameters:

layer_idx (int)

Return type:

None

class mstar.engine.resources.convenience.RaggedAttentionCallable(attn, label)[source]#

Bases: object

AttentionCallable for a cacheless attention: (q, k, v) -> out over one declared span of a RaggedAttnManager.

No kv and no layer cursor, because there is neither to carry: a ragged plan is keyed by the segment label, and the wrapper behind it is per (bucket, slot, label) rather than per layer. The label is therefore bound here at construction instead of being a cursor on the resource – which is what lets one stack attend over several spans at once (a refiner over the image tokens next to a pass over all of them), the case AttentionCallable says to thread explicitly.

One instance per label, held for the life of the resource binding: a compiled transformer region guards on the identity of the callables it is handed, so handing it a fresh one per step retraces that frame every step until Dynamo gives up.

Parameters: