mstar.model.vjepa2.submodules

Contents

mstar.model.vjepa2.submodules#

NodeSubmodule wrappers for the V-JEPA 2 graph nodes.

Three submodules:

VJepa2EncoderSubmodule - ViT 3D-patch video encoder. VJepa2PredictorSubmodule - masked latent predictor (single forward per call). VJepa2ACPredictorSubmodule - action-conditioned predictor (ditto).

All three are stateless (no KV cache, no per-iteration state) and are dispatched through StatelessEngine. preprocess stacks per-request tensors into a batch (dim 0) — the single-request case through _execute_sequential looks like B=1 and goes through forward; _execute_batched with B>1 goes through forward_batched.

The masked-vs-AC choice is made by VJepa2Model.get_submodule based on config.predictor_kind. Both predictor submodules emit the same output edge name (predicted_hidden) so downstream consumers don’t need to branch.

Classes

VJepa2ACPredictorSubmodule(predictor, config)

Action-conditioned predictor (V-JEPA 2-AC).

VJepa2ACRolloutPredictorSubmodule(predictor, ...)

Action-conditioned autoregressive rollout.

VJepa2EncoderSubmodule(encoder, config)

ViT 3D-patch video encoder.

VJepa2MPCPredictorSubmodule(predictor, config)

Single-request K-way AC predictor forward.

VJepa2MPCScorerSubmodule(config)

Score K candidate predicted latents against a goal latent.

VJepa2PredictorSubmodule(predictor, config)

Masked latent predictor (non-AC).

VJepa2RolloutPredictorSubmodule(predictor, ...)

Masked-predictor rollout submodule — autoregressive anticipation.
class mstar.model.vjepa2.submodules.VJepa2ACPredictorSubmodule(predictor, config)[source]#

Bases: ARNodeSubmodule

Action-conditioned predictor (V-JEPA 2-AC).

Additional required inputs vs the masked predictor: actions, states, and optionally extrinsics. Each is per-timestep, shape [T, action_embed_dim] (or [T, action_embed_dim - 1] for extrinsics), delivered as a GraphEdge by the model’s process_prompt / get_initial_forward_pass_args. Output is predicted_hidden with the same shape as encoder_hidden.

Parameters:
can_batch(batch, model_inputs)[source]#
Parameters:
Return type:

bool

forward(graph_walk, engine_inputs, encoder_hidden, actions, states, extrinsics=None, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, encoder_hidden, actions, states, extrinsics=None, **kwargs)[source]#

Batched form of forward: maps a multi-request batch to per-request outputs. Override when can_batch returns True.

Parameters:
Return type:

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

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

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor]

class mstar.model.vjepa2.submodules.VJepa2ACRolloutPredictorSubmodule(predictor, config)[source]#

Bases: ARNodeSubmodule

Action-conditioned autoregressive rollout.

Optionally uses a PiecewiseCudaGraphRunner to accelerate the inner block loop. The submodule opts in via get_piecewise_cuda_graph_configs (label "block_loop"); the engine builds the runner at warmup and passes it in through engine_inputs.piecewise_runners.

Sliding-window rollout — diverges from upstream vjepa2/notebooks/utils/mpc_utils.py::cem which uses growing-context (T: 1 → rollout+1) from a single-tubelet initial encoding. Our encoder’s natural output is T=grid_depth (a full 64-frame clip), so growing-context would immediately hit the AC predictor’s pre-built causal attn-mask cap, so we slide a fixed-size window instead.

Per iter k:

  • Slice actions/states: actions_k = actions[:, k : k + T_ctx] (T_ctx = config.grid_depth; client supplies the full trajectory once and we slice per-iter based on dynamic_loop_iter_counts).

  • Run VisionTransformerPredictorAC(encoder_hidden, actions_k, states_k) — returns [B, N, D] predictions at every one of the T_ctx timesteps.

  • Take the last tubelet group (predicted[:, -grid² :] — shape [B, grid², D]) as the “new” imagined state, matching upstream world_model’s next_frame = predictor_output[:, -HW:, :] interpretation.

  • Slide the encoder_hidden window: drop oldest grid² tokens, append the new grid² tokens. Fixed shape [B, N, D] every iter — torch.compile-friendly, no dynamic-shape recompiles.

actions / states (and optionally extrinsics) are returned unchanged as identity loop-back edges so the graph dispatcher can keep routing them on every iter. Per-iter slicing happens inside the forward based on iter_idx — the tensors themselves don’t change.

Shares the sibling rollout submodule’s batching contract (B >= 2 + shape + same-iter homogeneity) and check_stop early-exit semantics.

Parameters:
can_batch(batch, model_inputs)[source]#

Same rule as masked rollout: B >= 2, shape-homogeneous across encoder_hidden / actions / states (+ optional extrinsics), and same iter_idx for every rid.

B=1 → sequential forward path (see VJepa2EncoderSubmodule.can_batch for the AC warm-latency regression that motivates this gate).

Parameters:
Return type:

bool

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

Return the set of dynamic-loop names that should stop after this step.

Runs on the worker’s slow-postprocess path after execute_batch returns — never inside execute_batch. Allowed to read tensor values (.item() / .cpu()) because by this point the GPU thread is no longer blocked by it.

Stops returned here are deferred by one step: they apply to the worker’s next iter’s fast postprocess. The current in-flight step (already submitted under the assumption that the rid continues) will run for that rid and its output discarded — the standard 1-wasted-step cost for any stop signal.

Default: no stops.

Parameters:
Return type:

set[str]

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

The rollout step’s KV, when the block loop runs inline.

The region declares its own work whenever it holds a slot, so this only covers the eager path — which had no declaration at all: nothing reserved its pages, so running out of KV surfaced as a RuntimeError from inside the forward instead of an admit failure the worker can push back and evict for.

Parameters:
Return type:

SubmoduleStep | None

forward(graph_walk, engine_inputs, encoder_hidden, actions, states, extrinsics=None, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, encoder_hidden, actions, states, extrinsics=None, **kwargs)[source]#

Batched form of forward: maps a multi-request batch to per-request outputs. Override when can_batch returns True.

Parameters:
Return type:

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

get_piecewise_cuda_graph_configs(device, autocast_dtype, tp_world_size=1, **kwargs)[source]#

One BATCHED piecewise graph ("block_loop") for the AC predictor.

capture_seq_len = cond_tokens + N*N is the per-frame token count (one rollout step processes one frame). The block loop reads the KV cache, so the region declares its own KVStep / AttentionStep and takes its slot with lease_before_step=True.

Parameters:
Return type:

dict[str, PiecewiseCudaGraphConfig]

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

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor]

class mstar.model.vjepa2.submodules.VJepa2EncoderSubmodule(encoder, config)[source]#

Bases: NodeSubmodule

ViT 3D-patch video encoder.

Supports cross-request batching: requests with identical video_frames shapes are stacked on dim 0 and run in a single encoder forward. Shape-heterogeneous batches fall through to the engine’s sequential path (one forward per request).

Because VJepa2Model.process_prompt always resizes + center-crops to (crop_size, crop_size) and uniformly samples to frames_per_clip, concurrent requests hitting the same serving config are shape-homogeneous by construction. Cross-config deployments (rare) naturally fall back.

Parameters:
can_batch(batch, model_inputs)[source]#
Parameters:
forward(graph_walk, engine_inputs, pixel_values_videos, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

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

Batched encoder forward. Returns per-rid outputs directly.

The returned per-rid encoder_hidden slices preserve the leading batch dim as size 1 (hidden[i:i+1]) so downstream submodules see [1, N, D] — identical shape to what the sequential path emits. That keeps the preprocess stacking symmetric across paths and avoids a shape discrepancy at the predictor boundary.

Parameters:
Return type:

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

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

NodeInputs

preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor]

class mstar.model.vjepa2.submodules.VJepa2MPCPredictorSubmodule(predictor, config)[source]#

Bases: ARNodeSubmodule

Single-request K-way AC predictor forward.

Inputs:

  • encoder_hidden: [1, N, D] — context-video latent from the preceding video_encoder node.

  • actions: [K, T, 7] — K candidate action trajectories.

  • states: [K, T, 7] — matching proprioceptive states.

  • extrinsics (optional, use_extrinsics=True on the config).

Output: predicted_hidden of shape [K, N, out_dim] — all K candidates’ predicted future latents.

Parity anchor: vjepa2/notebooks/utils/mpc_utils.py::cem lines 62-64 (context expansion) + line 114 (world_model call with K-batched inputs).

Parameters:
forward(graph_walk, engine_inputs, encoder_hidden, actions, states, extrinsics=None, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

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

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor]

class mstar.model.vjepa2.submodules.VJepa2MPCScorerSubmodule(config)[source]#

Bases: NodeSubmodule

Score K candidate predicted latents against a goal latent.

Inputs:

  • predicted_hidden: [K, N, D] — from the MPC predictor.

  • goal_hidden: [1, N, D] (or [N, D]) — pre-encoded by the client via a prior prefill_video_encoder_only call.

Outputs emitted to the client:

  • best_index (int64 scalar): argmin of costs.

  • costs ([K]): per-candidate cost values.

  • predicted_hidden ([K, N, D]): passed through so clients that want the imagined trajectories (e.g. for visualization) get them without a second request.

Cost function is selectable via config.mpc_cost_fn:

  • "l1" (default): (pred - goal).abs().mean(dim=[1, 2]) — matches upstream mpc_utils.py::l1.

  • "l2": squared-error mean.

  • "cosine": 1 - cosine similarity (so argmin still picks best).

Parameters:

config (VJepa2Config)

forward(graph_walk, engine_inputs, predicted_hidden, goal_hidden, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

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

NodeInputs

preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor]

class mstar.model.vjepa2.submodules.VJepa2PredictorSubmodule(predictor, config)[source]#

Bases: ARNodeSubmodule

Masked latent predictor (non-AC).

Expects encoder_hidden from the preceding encoder node. Optionally accepts pre-built context_mask and target_mask edges; when absent, falls back to full-context / full-target defaults. Output is the predicted hidden tensor at the target positions.

Supports cross-request batching when encoder_hidden shapes agree across the batch AND (if provided) mask shapes agree. All-defaults is the common case — the full-context / full-target masks are derived from encoder_hidden.shape and so are trivially homogeneous when the encoder outputs are.

Parameters:
can_batch(batch, model_inputs)[source]#
Parameters:
Return type:

bool

forward(graph_walk, engine_inputs, encoder_hidden, context_mask=None, target_mask=None, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, encoder_hidden, context_mask=None, target_mask=None, **kwargs)[source]#

Batched form of forward: maps a multi-request batch to per-request outputs. Override when can_batch returns True.

Parameters:
Return type:

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

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

parts: list[torch.Tensor] = []

for inp in per_request_inputs:

t = inp[field_name][0] parts.append(_ensure_lead_batch_dim(t, target_rank))

return torch.cat(parts, dim=0)

Parameters:
Return type:

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor]

class mstar.model.vjepa2.submodules.VJepa2RolloutPredictorSubmodule(predictor, config, num_output_frames=2, frames_per_second=4, anticipation_seconds=1.0)[source]#

Bases: ARNodeSubmodule

Masked-predictor rollout submodule — autoregressive anticipation.

Parity target: upstream vjepa2/evals/action_anticipation_frozen/modelcustom/vit_encoder_predictor_concat_ar.py AnticipativeWrapper.forward (lines 172-224). That Python loop calls the (stateless) predictor once per rollout step, feeding a sliding window of the previous step’s prediction back as context:

x_pred_input = x_full # initial encoder output for _ in range(num_steps):

x_pred = predictor(x_pred_input, ctxt_positions, tgt_positions) x_pred_input = cat([x_pred_input[:, N_pred:], x_pred], dim=1) x_accumulate = cat([x_accumulate, x_pred], dim=1)

In mstar this becomes a Loop whose section is a single node wrapping this submodule. On each iter:

  • The loop-back encoder_hidden carries the sliding window state. Iter 0 receives it from the preceding video_encoder node; iter k > 0 receives the previous iter’s emitted encoder_hidden.

  • predicted_hidden is emitted both as a (dangling) section output — which lets the Loop’s cache_outputs machinery pick it up — and into the Loop’s accumulated_outputs for client delivery.

The predictor nn.Module is the same instance used by VJepa2PredictorSubmodule — weights are loaded once and shared. Only the preprocess/forward rollout-awareness differs.

Uses the single-request sequential path (can_batch=False inherited from NodeSubmodule). Same cross-request batching story as the encoder — enabling it requires an engine fix.

Parameters:
can_batch(batch, model_inputs)[source]#

Batch only when every request is at the same rollout iter AND their encoder_hidden shapes agree.

Different iters can’t share a forward because the sliding-window math torch.cat([hidden[:, n_pred:], predicted], dim=1) is symmetric across the batch dim only if all rids moved through the same number of prior iterations. The scheduler groups same-iter requests together, so this is the common case; mixed-iter batches fall through to sequential.

B=1 → sequential forward; see VJepa2EncoderSubmodule.can_batch for the rationale (AC warm-latency regression when forward_batched is the only path). Rollout is especially sensitive because it calls forward_batched once per iter — amortizing a slow compile over H iters would multiply the pain.

Parameters:
Return type:

bool

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

Return the set of dynamic-loop names that should stop after this step.

Runs on the worker’s slow-postprocess path after execute_batch returns — never inside execute_batch. Allowed to read tensor values (.item() / .cpu()) because by this point the GPU thread is no longer blocked by it.

Stops returned here are deferred by one step: they apply to the worker’s next iter’s fast postprocess. The current in-flight step (already submitted under the assumption that the rid continues) will run for that rid and its output discarded — the standard 1-wasted-step cost for any stop signal.

Default: no stops.

Parameters:
Return type:

set[str]

forward(graph_walk, engine_inputs, encoder_hidden, **kwargs)[source]#

One rollout step. Runs the masked predictor, then slides the encoder_hidden window forward so iter k+1’s context is the concatenation of the last (N - N_pred) context tokens plus the N_pred newly-predicted tokens.

Parameters:
Return type:

dict[str, list[Tensor]]

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

Batched rollout step across shape-homogeneous same-iter requests.

can_batch guarantees all rids are at the same iter_idx (else we fall to sequential), so the rollout-step math is a single [B, N, D] forward. Per-request early-exit (check_stop) still fires independently when each rid’s rollout_horizon is reached — individual rids can drop out while others continue, and the scheduler re-batches remaining rids on the next iter.

Parameters:
Return type:

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

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

ARNodeInputs

preprocess(graph_walk, engine_inputs, inputs)[source]#
Parameters:
Return type:

dict[str, Tensor]

Contents