mstar.model.waypoint.components

mstar.model.waypoint.components#

Waypoint-1.5-1B component modules.

The DiT is a native port of world_engine/src/model/world_model.py. apply_inference_patches runs unconditionally in the reference, so the patched model is the shipped one – but the port does not follow it uniformly: it takes the patched fused QKV and keeps the unpatched packed MLPFusion.fc1, since that patch splits a weight the checkpoint stores packed. The two forms are algebraically identical.

Nothing here owns the world state. The ring KV cache and the FlexAttention kernel are engine resources (engine/resources/kv/ring/, engine/resources/attn/flex.py) bound onto WaypointDiT and its 24 WaypointAttention layers at load by NodeSubmodule.bind_node_resources. That keeps the rings out of the module tree, where to_empty(device), state_dict() and the weight loader would each have a buffer of ours to leave holding garbage.

taehv.py is the exception: a streaming AE session is per-request Python state, held in PerRequestState by the two VAE nodes and dropped with the request. Its taehv imports are deferred, so this package imports without the upstream package installed.

class mstar.model.waypoint.components.AdaLN(dim)[source]#

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)[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:

x (Tensor)
cond (Tensor)

Return type:

Tensor

class mstar.model.waypoint.components.ChunkedStreamingTAEHV(ae_model, auto_aspect_ratio=True, device=None, dtype=torch.bfloat16, height=None, width=None)[source]#

Bases: object

One streaming session: this request’s frames, in order.

The order-dependent state lives here; ae_model is shared and never mutated, so a fresh instance is the reference’s reset(). Not an nn.Module: the module tree would give state_dict(), to() and the weight loader a second path to the same parameters.

Parameters:

ae_model (Module)
auto_aspect_ratio (bool)
device (device | None)
dtype (dtype)
height (int | None)
width (int | None)

decode(latent)[source]#

One latent [B, C, h, w] -> [t_upscale, H, W, 3] uint8.

Single-use and order-dependent: the temporal memory advances per call, so a latent decoded twice, out of turn or not at all shifts every frame after it, silently. The first call primes that memory with frames_to_trim extra feeds, which is why the prime walk decodes.

Parameters:: latent (Tensor)
Return type:: Tensor

encode(frames)[source]#

[t_downscale, H, W, 3] in [0, 1] -> one latent [B, C, h, w]. Exactly that many frames: fewer buffer and return None. The reference scales from uint8 here; this takes it pre-scaled in the same cast-then-divide order, so the two agree bit for bit.

Parameters:: frames (Tensor)
Return type:: Tensor

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

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)[source]#

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)[source]#

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.ControllerInputEmbedding(config)[source]#

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)[source]#

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.DeviceTableCache(*tables)[source]#

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)[source]#

Parameters:: device (device)
Return type:: tuple[Tensor, …]

class mstar.model.waypoint.components.MLP(dim_in, dim_middle, dim_out)[source]#

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)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Parameters:: x (Tensor)
Return type:: Tensor

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

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)[source]#

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.NoiseConditioner(dim, fourier_dim=512, base=10_000.0, *, reference_compat=False, cached_sigmas=())[source]#

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)[source]#

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)[source]#

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

Parameters:: device (device | str)
Return type:: None

class mstar.model.waypoint.components.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

Parameters:

x (Tensor)
rope_angles (tuple[Tensor, Tensor])

Return type:

Tensor

class mstar.model.waypoint.components.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:

x_pos (Tensor)
y_pos (Tensor)
t_pos (Tensor)

Return type:

tuple[Tensor, Tensor]

materialize(device)[source]#

Copy derived frequency tables before compilation/capture warmup.

Parameters:: device (device | str)
Return type:: None

class mstar.model.waypoint.components.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:

config (WaypointConfig)
layer_idx (int)

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:

x (Tensor)
frame_pos (Tensor)
rope_angles (tuple[Tensor, Tensor])
v1 (Tensor | None)
commit (bool)

Return type:

tuple[Tensor, Tensor]

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

Bases: Module

Conv2d patchify -> 24 blocks -> adaLN head -> unpatchify, plus the 4+1 per-frame driver.

Built on the meta device and materialized by weight_loader in a fixed order:

with torch.device("meta"):
    dit = WaypointDiT(config)
dit.cast_serving_dtypes()   # bf16 + fp32 islands, on meta
dit.to_empty(device=device) # un-ties cond_proj
dit.retie_cond_proj()       # MUST follow to_empty
load_weights_into(dit, ...)

The KV ring is derived state owned by an engine resource, bound after materialization – not part of this module.

Parameters:: config (WaypointConfig)

append_frame(latent, frame_pos, *, mouse, button, scroll)[source]#

Prime the world state from a real (VAE-encoded) frame: the committing pass alone, one forward.

latent is already the settled x0, so there is nothing to solve and nothing to clone. Returned unchanged for the caller to decode.

Parameters:

latent (Tensor)
frame_pos (Tensor)
mouse (Tensor)
button (Tensor)
scroll (Tensor)

Return type:

Tensor

cast_serving_dtypes()[source]#

bf16 everywhere, then the fp32 islands back to fp32.

Called on the meta module before to_empty(device) so storage is allocated directly in the serving dtype. .to(dtype) on meta preserves the cond_proj aliasing (to_empty does not), so no retie is needed here.

Return type:: WaypointDiT

compile_regions(**compile_kwargs)[source]#

Compile the denoise and cache passes, one region each. Idempotent.

materialize_runtime_tables must run first: the derived token grid, RoPE tables and conditioner LUT live outside the module tree, and a first touch inside a fullgraph=True region is at best a specialization and at worst a graph break.

Return type:: WaypointDiT

property dtype: dtype#: Bulk compute dtype (the non-island weights); callers cast latents to this.

forward(x, sigma, frame_pos, *, mouse, button, scroll, commit, cond_idx)[source]#

One pass over one latent frame; returns the rectified-flow velocity.

x [B, N, C, H, W] (N == 1), sigma [B, N], frame_pos [B] int64 ring clock. Returns [B, N, C, H, W]. commit says whether this pass keeps its K/V (False for the four denoise passes, True for the fifth); cond_idx selects the cached modulation row and must match sigma.

Parameters:

x (Tensor)
sigma (Tensor)
frame_pos (Tensor)
mouse (Tensor)
button (Tensor)
scroll (Tensor)
commit (bool)
cond_idx (int)

Return type:

Tensor

generate_frame(noise, frame_pos, *, mouse, button, scroll)[source]#

Denoise one frame from noise [B, N, C, H, W] and commit it.

Five forwards: 4 frozen + 1 committing, all sharing one frame_pos. The caller owns the ring clock and advances it once per committed frame.

Parameters:

noise (Tensor)
frame_pos (Tensor)
mouse (Tensor)
button (Tensor)
scroll (Tensor)

Return type:

Tensor

materialize_runtime_tables(device)[source]#

Create every derived table after weight loading and before compile.

Parameters:: device (device | str)
Return type:: WaypointDiT

retie_cond_proj()[source]#

Alias blocks 1..23’s six cond_proj matrices onto block 0’s.

Public, and separate from __init__, because to_empty(device) destroys the tying (Module._apply has no cross-module memo, so each block would get 24 independent copies). bias_in is deliberately NOT tied – it is genuinely per-layer.

Return type:: WaypointDiT

class mstar.model.waypoint.components.WaypointDiTBlock(config, layer_idx)[source]#

Bases: Module

One DiT block: adaLN-modulated causal frame attention, optional controller fusion, adaLN-modulated MLP.

The six modulation tensors come from this block’s cond_head, whose bias_in is per-layer and whose cond_proj matrices alias block 0’s.

Parameters:

config (WaypointConfig)
layer_idx (int)

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

x [B, N*T, D], cond/ctrl_emb [B, N, D] (per frame) -> (x, v1). v1 is layer 0’s pre-lerp V, threaded down the stack; cond_idx is the scheduler_sigmas slot for the cond_head cache. Only f_pos reaches this far – RoPE angles are built once at the root.

Parameters:

x (Tensor)
frame_pos (Tensor)
rope_angles (tuple[Tensor, Tensor])
cond (Tensor)
ctrl_emb (Tensor)
v1 (Tensor | None)
commit (bool)
cond_idx (int)

Return type:

tuple[Tensor, Tensor]

class mstar.model.waypoint.components.WaypointPosIds(f_pos, t_pos, y_pos, x_pos)[source]#

Bases: NamedTuple

The four position streams for one frame, per row of the step batch.

f_pos is the [B] int64 ring clock (the KV resource’s upsert takes it); t_pos is the [B, T] RoPE time coordinate f_pos * ts_mult; y_pos/x_pos are [B, T] token-grid coordinates, row = i // width and col = i % width.

Parameters:

f_pos (Tensor)
t_pos (Tensor)
y_pos (Tensor)
x_pos (Tensor)

f_pos: Tensor#: Alias for field number 0

t_pos: Tensor#: Alias for field number 1

x_pos: Tensor#: Alias for field number 3

y_pos: Tensor#: Alias for field number 2

mstar.model.waypoint.components.ada_gate(x, gate)[source]#

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.ada_rmsnorm(x, scale, bias)[source]#

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

x (Tensor)
rope_angles (tuple[Tensor, Tensor])

Return type:

Tensor

mstar.model.waypoint.components.load_taehv(ae_uri, cache_dir=None)[source]#

Load the shared weights module from a resolved local checkpoint.

Hub resolution belongs to checkpoint.resolve_taehv_checkpoint, so only the required file is downloaded. cache_dir is accepted for compatibility with existing callers but is unused here.

Parameters:

ae_uri (str)
cache_dir (str | None)

Return type:

Module

mstar.model.waypoint.components.rms_norm(x)[source]#

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

Modules

`attention`	Waypoint self-attention: fused QKV, value residual, OrthoRoPE, ring cache.
`dit`	The Waypoint-1.5 DiT: 24 blocks, the 4+1 pass driver, and the world clock.
`layers`	Stateless layer primitives for the Waypoint-1.5 DiT.
`rope`	OrthoRoPE: Waypoint's orthogonal (x, y, t) rotary position embedding.
`taehv`

mstar.model.waypoint.components

Contents

mstar.model.waypoint.components#