mstar.model.waypoint.components.attention

mstar.model.waypoint.components.attention#

Waypoint self-attention: fused QKV, value residual, OrthoRoPE, ring cache.

The served reference is the fused patch_model.MergedQKVAttn, not three separate GEMMs, so that fused form is the parity target. qkv_proj.weight is cat([q, k, v], dim=0); GQA makes the K/V slabs unequal, so swapping them loads cleanly but serves wrong video.

Classes

WaypointAttention(config, layer_idx)

One layer of causal frame attention over the ring KV cache.
class mstar.model.waypoint.components.attention.WaypointAttention(config, layer_idx)[source]#

Bases: Module

One layer of causal frame attention over the ring KV cache.

x is [B, N*T, D] (N frames of T tokens, flattened; N is 1 for every served call) and the head layout inside is [B, H, N*T, d_head], 32 query heads over 16 KV heads.

Parameters:
bind_resources(resources)[source]#

Resolve the two resources this layer calls.

NodeSubmodule.bind_node_resources walks self.modules(), so one bind on the submodule reaches all 24 layers. .get: a layer may be bound on a node that owns only some of them.

Parameters:

resources (dict)

Return type:

None

forward(x, frame_pos, rope_angles, v1, *, commit)[source]#

x [B, N*T, D] -> (out [B, N*T, D], v1 [B, H_kv, N*T, d_head]).

frame_pos is the [B] int64 ring clock; v1 is None at layer 0 and layer 0’s pre-lerp V thereafter. Both stay arguments rather than cursors the KV resource keeps: a cursor that drifts from the caller does not raise, it rewrites history.

Parameters:
Return type:

tuple[Tensor, Tensor]