mstar.model.waypoint.submodules

mstar.model.waypoint.submodules#

Classes

WaypointDitSubmodule(dit, taehv, config)

The world DiT: one latent frame per engine step, TAEHV-decoded in the same forward.

WaypointVaeEncoderSubmodule(taehv, config)

TAEHV encoder: temporal_compression raw frames -> one latent frame.
class mstar.model.waypoint.submodules.WaypointDitSubmodule(dit, taehv, config)[source]#

Bases: _FunctionalAeMixin, NodeSubmodule

The world DiT: one latent frame per engine step, TAEHV-decoded in the same forward.

The decode is fused in rather than left on its own node so a same-worker speculative N+1 (GraphNode.enable_async_scheduling) can start while N’s frame is still going out. See WaypointModel’s module docstring for the resulting two-node graph.

Parameters:
bind_node_resources(resources)[source]#

Require both resources before letting the bind reach the layers.

Parameters:

resources (dict)

Return type:

None

can_batch(batch, model_inputs)[source]#

Rows are independent (preprocess concatenates on the batch dim and every resource is scoped per row), so any admitted set is batchable.

Return type:

bool

check_stop(request_id, request_info, outputs)[source]#

Stop the rollout loop after exactly num_steps iterations.

Iteration k (0-based) is still being postprocessed when the loop counter reads k, so N frames means firing at k + 1 >= N. Under async scheduling an overshoot iteration this signal is too late to stop is vetoed instead in prepare_inputs.

Parameters:
Return type:

set[str]

cleanup_request(request_id)[source]#

Drop the request’s clock. The ring itself is released by the engine’s remove_request sweep (RingKVManager.remove_request), not here.

Parameters:

request_id (str)

declare_step(graph_walk, request_ids, inputs, slot_lease=None, piecewise_leases=None, **kwargs)[source]#

Name both resources so the runner drives their lifecycle.

Neither step carries segments or reserves anything: a ring overwrites in place, so there is no span to admit. Declaring them is still required for admit to hand a request one of the node’s worlds. The KV step carries the ring clock via RingKVStep, letting RingKVManager.admit check it against the frame its last commit recorded for that rid.

Parameters:
Return type:

SubmoduleStep

disable_autocast: bool = True#
disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, frame_pos, mouse, button, scroll, noise=None, latent=None, **kwargs)[source]#

One frame. rollout denoises it from noise, prime appends a real one; both commit to the ring, decode it with TAEHV in the same call, and return the decoded frame, the updated histories, and the clock passthrough.

engine_inputs is unused here, on purpose: under capture it is the dummy request’s forever. The ring and attention backend come off self.node_resources, resolved once at bind_node_resources.

Parameters:
Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, **kwargs)[source]#

Run the batched forward once, then split its rows back out.

Row i of every output tensor belongs to engine_inputs.request_ids[i], the same row-order invariant preprocess used going in.

Parameters:
Return type:

dict[str, dict[str, list[Tensor]]]

get_accelerator_graph_configs(device, tp_world_size=1)[source]#

Both walks, as optional captures.

Prime’s inputs are the rollout template with noise renamed, so it reuses the pool the rollout graph already sized.

Parameters:
Return type:

list[AcceleratorGraphConfig]

max_batch_size(graph_walk)[source]#

Both walks carry up to step_batch_size rows: one per resident world sharing the forward.

Parameters:

graph_walk (str)

Return type:

int

postprocess(request_id, request_info, outputs, inputs=None, **kwargs)[source]#

Advance the ring clock by exactly one committed frame, and copy the fused decode’s updated histories into the request’s stable tensors.

Both walks commit, so both advance the clock. The clock advance is metadata only, and the history copy is a device copy_ into the same fixed-address tensors prepare_inputs reads back next call – neither syncs.

Parameters:
prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#

This frame’s row: the ring clock, its controller slice, and either the noise to denoise from (rollout) or the latent to prime with. Also the nine decoder histories the fused decode reads and writes.

Runs on the host, outside any captured region, which is why the noise is drawn here rather than in forward: a captured region cannot call the RNG.

inputs.get("clock") is never read; it only has to be present in input_names for GraphNode.is_ready_for_speculation to propose the next iteration as a same-node speculation target.

Parameters:
Return type:

NodeInputs | None

preprocess(graph_walk, engine_inputs, inputs)[source]#

Concatenate every row’s tensors along the batch dim, in engine_inputs.request_ids order (inputs[i] pairs positionally with row i – the row-order invariant every batched resource below this node relies on). len(inputs) == 1 returns the row unchanged, no copy.

Parameters:
Return type:

dict

class mstar.model.waypoint.submodules.WaypointVaeEncoderSubmodule(taehv, config)[source]#

Bases: _SingleRequestMixin, _FunctionalAeMixin, NodeSubmodule

TAEHV encoder: temporal_compression raw frames -> one latent frame.

image_inputs [4, 720, 1280, 3] uint8 -> latent [1, 1, 32, 32, 64], the dit’s priming input unchanged. Prime walk only.

Parameters:
cg_key_info(graph_walk, per_request_info, per_request_input_metadata=None, **kwargs)[source]#

The captured (H, W) bucket for this batch’s seed, or None for eager.

The seed shape isn’t on model_inputs at lease time, so it is read off step_metadata["seed_hw"], set by get_initial_forward_pass_args.

Parameters:
Return type:

tuple[int, int] | None

disable_autocast: bool = True#
disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, image, **kwargs)[source]#

Encode the seed clip into the latent the dit primes the world on. Request ids are never read here: forward deletes engine_inputs outright, so it never needs a real request’s ids, captured or not.

Parameters:
Return type:

dict[str, list[Tensor]]

get_accelerator_graph_configs(device, tp_world_size=1)[source]#
Parameters:
Return type:

list[AcceleratorGraphConfig]

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#

The seed clip, scaled into the AE’s dtype. Cast then divide, the reference’s order: 0-255 is exact in bf16, so the divide rounds once.

Parameters:
Return type:

NodeInputs