mstar.model.waypoint.components.layers

Stateless layer primitives for the Waypoint-1.5 DiT.

Port of world_engine/src/model/nn.py plus the three conditioning modules from world_engine/src/model/world_model.py. Parameter names track checkpoint keys (fc1/fc2, bias_in, cond_proj, mlp), not mstar’s usual components.mlp.MLP spelling, so this doesn’t reuse it.

Functions

ada_gate(x, gate)	Per-frame adaLN output gate: x * gate, same shapes as ada_rmsnorm.
ada_rmsnorm(x, scale, bias)	Per-frame adaLN modulation: rms_norm(x) * (1 + scale) + bias.
bf16_roundtrip(table)	Quantize an fp32 derived table the way NoCastModule._apply does.
fp32_matmul_highest()	Run fp32 GEMMs as IEEE fp32, whatever the process-wide float32_matmul_precision (mstar.engine sets 'high').
rms_norm(x)	Unweighted RMS norm over the last dim (reference nn.rms_norm).

Classes

AdaLN(dim)	adaLN with the scale/shift projection folded in: one Linear produces [scale | shift] from silu(cond).
CondHead(config)	Per-layer WAN-style conditioning head: the noise embedding becomes the block's six adaLN modulation tensors (scale/shift/gate for attention, then the same three for the MLP).
ControllerInputEmbedding(config)	Controller state -> one conditioning vector per frame.
DeviceTableCache(*tables)	Per-device replicas of small derived fp32 tables (RoPE frequencies, Fourier frequencies), held outside the module tree so they stay out of state_dict/to_empty's reach and fp32 forever.
MLP(dim_in, dim_middle, dim_out)	Two-layer SiLU MLP, both projections bias-free.
MLPFusion(config)	Fuses a per-frame conditioning vector into that frame's tokens.
NoiseConditioner(dim[, fourier_dim, base, ...])	sigma -> Fourier features -> MLP, the DiT's only noise-level input.

class mstar.model.waypoint.components.layers.AdaLN(dim)

Bases: Module

adaLN with the scale/shift projection folded in: one Linear produces [scale | shift] from silu(cond).

The DiT’s output head only; blocks get their six modulation tensors from CondHead. cond is per-frame [B, N, D], expanded over its tokens.

Parameters:

dim (int)

forward(x, cond)

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:

• x (Tensor)

• cond (Tensor)

Return type:

Tensor

class mstar.model.waypoint.components.layers.CondHead(config)

Bases: Module

Per-layer WAN-style conditioning head: the noise embedding becomes the block’s six adaLN modulation tensors (scale/shift/gate for attention, then the same three for the MLP).

bias_in is genuinely per-layer; the six cond_proj Linears are physically shared across all 24 blocks via WaypointDiT.retie_cond_proj (must run after to_empty – see its docstring).

Parameters:

config (WaypointConfig)

build_cache(conds)

Precompute the six modulation tensors for every scheduled sigma.

conds is [S, 1, D] – one embedding per scheduler_sigmas entry, each built at M=1 exactly as forward receives it, so a cached row is bit-identical to the live _project it replaces. Switches forward to a gather; idempotent enough to rebuild.

Parameters:

conds (Tensor)

Return type:

None

forward(cond, cond_idx)

cond [B, N, D], cond_idx the scheduler_sigmas slot -> six [B, N, D] tensors (s0, b0, g0, s1, b1, g1).

After build_cache the six come from the cache row cond_idx and cond is unused

Parameters:

• cond (Tensor)

• cond_idx (int)

Return type:

tuple[Tensor, …]

n_cond = 6

class mstar.model.waypoint.components.layers.ControllerInputEmbedding(config)

Bases: Module

Controller state -> one conditioning vector per frame.

The concat order is (mouse, button, scroll) – 2 + n_buttons + 1 = 259 here – and it is a checkpoint fact. A permuted order does not raise: the widths still sum to 259 and every downstream shape checks out, it just reads mouse velocity out of button columns. Callers pass the three positionally.

Parameters:

config (WaypointConfig)

forward(mouse, button, scroll)

mouse [B, N, 2], button [B, N, n_buttons], scroll [B, N, 1] -> [B, N, d_model].

Parameters:

• mouse (Tensor)

• button (Tensor)

• scroll (Tensor)

Return type:

Tensor

class mstar.model.waypoint.components.layers.DeviceTableCache(*tables)

Bases: object

Per-device replicas of small derived fp32 tables (RoPE frequencies, Fourier frequencies), held outside the module tree so they stay out of state_dict/to_empty’s reach and fp32 forever.

Callers MUST build the CPU tables with an explicit device="cpu": __init__ runs under with torch.device("meta"), where an ambient-device torch.arange produces a data-less meta tensor.

Parameters:

tables (Tensor)

get(device)

Parameters:

device (device)

Return type:

tuple[Tensor, …]

class mstar.model.waypoint.components.layers.MLP(dim_in, dim_middle, dim_out)

Bases: Module

Two-layer SiLU MLP, both projections bias-free.

bias=False is a checkpoint fact: the reference’s MLPFusion calls F.linear(h, self.mlp.fc2.weight) with no bias argument, which is only correct because there is no bias to pass.

Parameters:

• dim_in (int)

• dim_middle (int)

• dim_out (int)

forward(x)

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:

x (Tensor)

Return type:

Tensor

class mstar.model.waypoint.components.layers.MLPFusion(config)

Bases: Module

Fuses a per-frame conditioning vector into that frame’s tokens.

Parameter tree is MLP(2*D, D, D) over cat([x, cond]) – one [D, 2D] mlp.fc1 for the loader to fill. Compute splits that matrix instead (chunk(2, dim=1)) so cond broadcasts over the T tokens of its frame, compute-time only – stored split parameters would undo the loader’s concat into mlp.fc1.

Parameters:

config (WaypointConfig)

forward(x, cond)

x [B, N*T, D], cond [B, N, D] -> [B, N*T, D].

Parameters:

• x (Tensor)

• cond (Tensor)

Return type:

Tensor

class mstar.model.waypoint.components.layers.NoiseConditioner(dim, fourier_dim=512, base=10_000.0, *, reference_compat=False, cached_sigmas=())

Bases: Module

sigma -> Fourier features -> MLP, the DiT’s only noise-level input.

fp32 island: FP32_MODULE_PATHS pins self.mlp back to fp32 after the serving bf16 cast, and the body runs under autocast(enabled=False) on a .float() sigma and, eagerly, under fp32_matmul_highest (a batched call rounds through TF32 at 'high', and cuBLAS 13 on a B200 rounds even the served [1, 512] @ [512, 8192] GEMV). Compiled regions read the rows WaypointDiT builds eagerly instead of running the island. * 1000 scales [0, 1] into the Fourier basis’s rotating range; * 2**0.5 restores unit variance after the sin/cos concat.

Parameters:

• dim (int)

• fourier_dim (int)

• base (float)

• reference_compat (bool)

• cached_sigmas (tuple[float, ...])

forward(s)

s is [B, N] sigma; returns [B, N, dim] in s’s dtype (bf16 under reference_compat, whose table is stored bf16).

Parameters:

s (Tensor)

Return type:

Tensor

materialize(device)

Materialize frequency data and the optional post-load sigma LUT.

Parameters:

device (device | str)

Return type:

None

mstar.model.waypoint.components.layers.ada_gate(x, gate)

Per-frame adaLN output gate: x * gate, same shapes as ada_rmsnorm. No 1 + here – the gate multiplies the sublayer output before the residual add, so a zero gate contributes nothing.

Parameters:

• x (Tensor)

• gate (Tensor)

Return type:

Tensor

mstar.model.waypoint.components.layers.ada_rmsnorm(x, scale, bias)

Per-frame adaLN modulation: rms_norm(x) * (1 + scale) + bias.

x is [B, N*T, D]; scale and bias are [B, N, D], one modulation vector per frame. The unflatten is what makes those vectors broadcast over their own frame’s tokens.

Parameters:

• x (Tensor)

• scale (Tensor)

• bias (Tensor)

Return type:

Tensor

mstar.model.waypoint.components.layers.bf16_roundtrip(table)

Quantize an fp32 derived table the way NoCastModule._apply does. Only reachable under WaypointConfig.reference_compat.

Parameters:

table (Tensor)

Return type:

Tensor

mstar.model.waypoint.components.layers.fp32_matmul_highest()

Run fp32 GEMMs as IEEE fp32, whatever the process-wide float32_matmul_precision (mstar.engine sets 'high'). Eager only: the setting is global, so dynamo cannot trace a switch of it.

mstar.model.waypoint.components.layers.rms_norm(x)

Unweighted RMS norm over the last dim (reference nn.rms_norm).

No learned gain: every use site either gets its scale from adaLN or wants a bare normalization (Q/K norm, the two MLPFusion inputs).

Parameters:

x (Tensor)

Return type:

Tensor