mstar.model.waypoint.components.rope

mstar.model.waypoint.components.rope#

OrthoRoPE: Waypoint’s orthogonal (x, y, t) rotary position embedding.

d_xy and d_t count rotation pairs, not dims: at d_head == 64 x owns dims 0..15, y 16..31 and t 32..63, disjoint bands whose phases add independently. The cache stores post-RoPE keys, so replayed history is never re-rotated.

Functions

apply_ortho_rope(x, rope_angles)

Rotate x [B, H, T, d_head] by the (cos, sin) tables.

Classes

OrthoRoPE([config])

Stateless module wrapper around apply_ortho_rope(), constructed per attention layer to match the reference's call site.

OrthoRoPEAngles(config)

Builds the shared (cos, sin) angle tables for one forward, once under the DiT rather than once per block.
class mstar.model.waypoint.components.rope.OrthoRoPE(config=None)[source]#

Bases: Module

Stateless module wrapper around apply_ortho_rope(), constructed per attention layer to match the reference’s call site. config is stored for that parity and otherwise unused.

Parameters:

config (WaypointConfig | None)

forward(x, rope_angles)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

class mstar.model.waypoint.components.rope.OrthoRoPEAngles(config)[source]#

Bases: Module

Builds the shared (cos, sin) angle tables for one forward, once under the DiT rather than once per block.

Parameters:

config (WaypointConfig)

forward(x_pos, y_pos, t_pos)[source]#

[B, T] integer position grids -> (cos, sin), each [B, 1, T, d_head // 2] fp32 with a broadcast head axis.

t_pos is the RoPE clock, not the ring clock f_pos – equal here (ts_mult == 1), but they diverge at other serving fps.

Parameters:
Return type:

tuple[Tensor, Tensor]

materialize(device)[source]#

Copy derived frequency tables before compilation/capture warmup.

Parameters:

device (device | str)

Return type:

None

mstar.model.waypoint.components.rope.apply_ortho_rope(x, rope_angles)[source]#

Rotate x [B, H, T, d_head] by the (cos, sin) tables.

The output is a permutation of the interleaved layout (unfold in, cat out); it cancels inside q @ k.T, and is kept so that stored K/Q compare bit-for-bit with the reference. fp32 island: x is upcast for the rotation and cast back after.

Parameters:
Return type:

Tensor