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

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