# Author: Rui Heng Yang """FasterWAMDecoupled -- decoupled variant with action-aligned video RoPE. This module defines a parallel variant where video self-attention still uses the video's native 3D RoPE, but the video K cached for action attention is kept without video-self-attention RoPE (except ``legacy_3d``, which caches post-3D-RoPE K). Before action attention concatenates video K/V with action K/V, the selected/fused video K is re-positioned in the action attention coordinate system. ``new_fused_kv`` can use aligned 3D RoPE (``aligned_3d``), the same plus per-head camera-region masking (``aligned_3dp``), temporal-only aligned 1D RoPE, original video 3D RoPE with action 1D RoPE, or ``legacy_3d`` (fuse post-3D-RoPE K, then no extra video RoPE like ``video_zero_1d``). """ from __future__ import annotations from typing import Any, Callable, Dict, Optional import torch from fastwam.utils.logging_config import get_logger from .action_dit import ActionDiT from .fastwam_decoupled import FastWAMDecoupled from .helpers.loader import load_wan22_ti2v_5b_components from .mot_decoupled import ACTION_EXPERT_KEY, VIDEO_EXPERT_KEY, MoTDecoupled from .wan_video_dit import modulate, precompute_freqs_cis, rope_apply logger = get_logger(__name__) def _validate_new_fused_kv_projection_mode( kv_source_mode: str, projection_mode: str, ) -> str: """Validate and normalize the projection-mode configuration.""" projection_mode = str(projection_mode) if projection_mode not in MoTDecoupled.supported_new_fused_kv_projection_modes: raise ValueError( f"Unknown new_fused_kv_projection_mode {projection_mode!r}; " "expected one of " f"{sorted(MoTDecoupled.supported_new_fused_kv_projection_modes)}" ) if kv_source_mode != "new_fused_kv" and projection_mode != "full": raise ValueError( "new_fused_kv_projection_mode is only configurable when " "kv_source_mode=new_fused_kv; use full for other modes." ) return projection_mode class MoTDecoupledActionAlignedVideoRoPE(MoTDecoupled): """Decoupled MoT with raw video K cache and action-aligned video K RoPE. Differences from ``MoTDecoupled``: - video self-attention uses native video 3D RoPE exactly as before; - cached video K for action attention is stored before 3D RoPE by default (``legacy_3d`` instead caches post-3D-RoPE K); - before concatenating ``[video K, action K]``, video K is re-positioned according to the selected KV/RoPE mode. """ enable_new_fused_kv_key_norm = True # EEF-relative camera RoPE siblings. They share every code path except # action-mask construction; see `_build_ee_rope_freqs`. ee_rope_modes = frozenset({"ee_rope", "exclusive_ee_rope"}) supported_new_fused_kv_rope_modes = { "aligned_3d", "aligned_3dp", "aligned_1d", "video_zero_1d", "original_3d", "legacy_3d", "ee_rope", "exclusive_ee_rope", } aligned_3d_action_spatial_anchor_presets = { "center": ((0.5, 0.5),), "horizontal": ((0.5, 0.25), (0.5, 0.75)), "libero": ((0.5, 0.25), (0.5, 0.75)), # LIBERO with the wrist anchor moved off the image centre and onto the # place the gripper actually appears. The wrist camera is rigidly mounted # to the end effector, so the EEF projects to the SAME camera-local token # in every frame of every task -- (5.358016, 2.993164) on the 7x7 # per-camera grid, measured across all 10030 LIBERO-Plus entries with # zero variance. `horizontal` puts the wrist head at composite # (3.0, 10.0); the true gripper is at composite (5.358016, 9.993164), so # the column was already right and only the row was off, by 2.36 tokens. # Normalized for `anchor * size - 0.5` on the 7x14 LIBERO composite: # row (5.358016 + 0.5) / 7, column (9.993164 + 0.5) / 14. # Reproduce by projecting the EEF origin through `wrist_mount_T_eef_from_C` # in configs/calibration/libero_plus_anchors_v1.json: the point lands at # the translation of inv(mount), independent of the robot's pose. "libero_wrist_grounded": ( (0.5, 0.25), (0.836859390551857, 0.74951171875), ), "robotwin": ( (1.0 / 3.0, 0.5), (5.0 / 6.0, 0.25), (5.0 / 6.0, 0.75), ), } _aligned_3d_family_rope_modes = frozenset({"aligned_3d", "aligned_3dp"}) # Geometry identity is a CONJUNCTION, not the calibration digest alone # (plan Section 20.1). A dataset re-exported at a different raw resolution # changes every anchor while leaving the digest untouched. EEF_GEOMETRY_IDENTITY_FIELDS = ( "calibration_digest", "raw_source_resolution", "token_grid_h", "token_grid_w", "camera_order", "projection_version", ) def __init__( self, *args, new_fused_kv_rope_mode: str = "aligned_3d", aligned_3d_action_spatial_anchor_layout: str | None = "center", eef_geometry_identity: dict[str, Any] | None = None, **kwargs, ): if new_fused_kv_rope_mode not in self.supported_new_fused_kv_rope_modes: raise ValueError( "Unknown new_fused_kv_rope_mode " f"{new_fused_kv_rope_mode!r}; expected one of " f"{sorted(self.supported_new_fused_kv_rope_modes)}" ) self.new_fused_kv_rope_mode = new_fused_kv_rope_mode anchor_layout = ( "center" if aligned_3d_action_spatial_anchor_layout is None else str(aligned_3d_action_spatial_anchor_layout) ) if anchor_layout not in self.aligned_3d_action_spatial_anchor_presets: raise ValueError( "Unknown aligned_3d_action_spatial_anchor_layout " f"{anchor_layout!r}; expected one of " f"{sorted(self.aligned_3d_action_spatial_anchor_presets)}" ) self.aligned_3d_action_spatial_anchor_layout = anchor_layout self.aligned_3d_action_spatial_anchors = ( self.aligned_3d_action_spatial_anchor_presets[anchor_layout] ) if eef_geometry_identity is not None: missing = [ field for field in self.EEF_GEOMETRY_IDENTITY_FIELDS if field not in eef_geometry_identity ] if missing: raise ValueError( "eef_geometry_identity is incomplete; geometry identity is a " f"conjunction and every field is load-bearing. Missing: {missing}" ) eef_geometry_identity = { field: eef_geometry_identity[field] for field in self.EEF_GEOMETRY_IDENTITY_FIELDS } elif new_fused_kv_rope_mode in self.ee_rope_modes: raise ValueError( f"new_fused_kv_rope_mode={new_fused_kv_rope_mode!r} requires " "eef_geometry_identity. Without it a checkpoint cannot record " "which geometry it was trained against, and a dataset " "re-exported at another raw resolution would silently produce " "different anchors under the same calibration digest " "(plan Section 20.1)." ) self.eef_geometry_identity = eef_geometry_identity super().__init__(*args, **kwargs) @property def _cache_post_3d_rope_video_k(self) -> bool: """Whether action-path video K should be cached after video 3D RoPE.""" return ( self.kv_source_mode == "new_fused_kv" and self.new_fused_kv_rope_mode == "legacy_3d" ) def _build_expert_attention_io_with_raw_k( self, expert, block, x: torch.Tensor, freqs: torch.Tensor, t_mod: torch.Tensor, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool, ]: """Build attention tensors and expose both raw and RoPE-applied K.""" shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self._split_modulation(block, t_mod) ) attn_input = modulate(block.norm1(x), shift_msa, scale_msa) q = block.self_attn.norm_q(block.self_attn.q(attn_input)) k_raw = block.self_attn.norm_k(block.self_attn.k(attn_input)) v = block.self_attn.v(attn_input) q = rope_apply(q, freqs, block.num_heads) k_rope = rope_apply(k_raw, freqs, block.num_heads) use_gradient_checkpointing = bool( getattr(expert, "use_gradient_checkpointing", False) ) return ( q, k_rope, k_raw, v, x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gradient_checkpointing, ) def _build_expert_attention_io_with_pre_norm_k( self, expert, block, x: torch.Tensor, freqs: torch.Tensor, t_mod: torch.Tensor, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool, ]: """Build video attention while exposing K before RMSNorm and RoPE.""" shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( self._split_modulation(block, t_mod) ) attn_input = modulate(block.norm1(x), shift_msa, scale_msa) q = block.self_attn.norm_q(block.self_attn.q(attn_input)) k_pre_norm = block.self_attn.k(attn_input) k_rope = rope_apply(block.self_attn.norm_k(k_pre_norm), freqs, block.num_heads) v = block.self_attn.v(attn_input) q = rope_apply(q, freqs, block.num_heads) use_gradient_checkpointing = bool( getattr(expert, "use_gradient_checkpointing", False) ) return ( q, k_rope, k_pre_norm, v, x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gradient_checkpointing, ) def _apply_action_zero_rope_to_video_k( self, k_video_raw: torch.Tensor, action_freqs: torch.Tensor, ) -> torch.Tensor: """Rotate every video key token with action token 0's 1D RoPE.""" if action_freqs.ndim != 3 or action_freqs.shape[0] < 1: raise ValueError( f"action_freqs must be [Sa, 1, rope_dim] with Sa >= 1, " f"got {tuple(action_freqs.shape)}" ) action_zero_freq = action_freqs[:1].to(device=k_video_raw.device) return rope_apply(k_video_raw, action_zero_freq, self.num_heads) def _split_3d_rope_dims(self) -> tuple[int, int, int]: """Return temporal/height/width RoPE complex dims for one attention head.""" temporal_real_dim = self.attn_head_dim - 2 * (self.attn_head_dim // 3) spatial_real_dim = self.attn_head_dim // 3 return ( temporal_real_dim // 2, spatial_real_dim // 2, spatial_real_dim // 2, ) def _infer_video_tokens_per_frame(self, video_freqs: torch.Tensor) -> int: """Infer the contiguous video tokens-per-frame from repeated temporal freqs.""" t_dim, _, _ = self._split_3d_rope_dims() if video_freqs.ndim != 3 or video_freqs.shape[0] < 1: raise ValueError( f"video_freqs must be [Sv, 1, rope_dim], got {tuple(video_freqs.shape)}" ) temporal = video_freqs[:, 0, :t_dim] first_temporal = temporal[:1] same_as_first = torch.isclose( temporal, first_temporal.expand_as(temporal), ).all(dim=-1) first_different = (~same_as_first).nonzero(as_tuple=False) if first_different.numel() == 0: return int(video_freqs.shape[0]) return int(first_different[0, 0].item()) def _infer_video_spatial_grid_size( self, video_freqs: torch.Tensor, tokens_per_frame: int, ) -> tuple[int, int]: """Infer one-frame token grid height/width from the 3D RoPE layout.""" if tokens_per_frame <= 1: return 1, max(tokens_per_frame, 1) t_dim, h_dim, w_dim = self._split_3d_rope_dims() first_frame = video_freqs[:tokens_per_frame, 0] w_start = t_dim + h_dim w_part = first_frame[:, w_start:w_start + w_dim] width = None for idx in range(1, tokens_per_frame + 1): if tokens_per_frame % idx != 0: continue if torch.allclose(w_part[idx - 1], w_part[-1]): width = idx break if width is None or width <= 0: width = tokens_per_frame return max(tokens_per_frame // width, 1), width def _infer_video_spatial_center_index( self, video_freqs: torch.Tensor, tokens_per_frame: int, ) -> int: """Infer the center token index inside one video frame from h/w RoPE cycles.""" if tokens_per_frame <= 1: return 0 t_dim, h_dim, w_dim = self._split_3d_rope_dims() first_frame = video_freqs[:tokens_per_frame, 0] w_start = t_dim + h_dim w_part = first_frame[:, w_start:w_start + w_dim] width = None for idx in range(1, tokens_per_frame + 1): if tokens_per_frame % idx != 0: continue if torch.allclose(w_part[idx - 1], w_part[-1]): width = idx break if width is None or width <= 0: return tokens_per_frame // 2 height = max(tokens_per_frame // width, 1) return min((height // 2) * width + (width // 2), tokens_per_frame - 1) def _build_2d_sincos_spatial_pe( self, height: int, width: int, dim: int, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: """Build fixed DETR-style 2D sine/cosine PE for one video frame.""" if dim % 4 != 0: raise ValueError(f"2D sin/cos PE dim must be divisible by 4, got {dim}") quarter = dim // 4 omega = torch.arange(quarter, device=device, dtype=torch.float32) omega = 1.0 / (10000 ** (omega / max(quarter, 1))) y = torch.arange(height, device=device, dtype=torch.float32) x = torch.arange(width, device=device, dtype=torch.float32) y_embed = y[:, None] * omega[None, :] x_embed = x[:, None] * omega[None, :] y_embed = torch.cat([y_embed.sin(), y_embed.cos()], dim=-1) x_embed = torch.cat([x_embed.sin(), x_embed.cos()], dim=-1) pe = torch.cat([ y_embed[:, None, :].expand(height, width, -1), x_embed[None, :, :].expand(height, width, -1), ], dim=-1) return pe.reshape(height * width, dim).to(dtype=dtype) def _build_simple_pe_for_video_k( self, k_video: torch.Tensor, action_layer_idx: int, video_freqs: torch.Tensor, ) -> torch.Tensor: """Build spatial-only fixed 2D sin/cos PE through the learned projection.""" projection = getattr(self, "k_video_pos_projection", None) if projection is None: raise ValueError( "simple+PE modes require k_video_pos_projection to be initialized." ) tokens_per_frame = self._infer_video_tokens_per_frame(video_freqs) height, width = self._infer_video_spatial_grid_size( video_freqs=video_freqs, tokens_per_frame=tokens_per_frame, ) if height * width != tokens_per_frame: raise ValueError( f"Invalid inferred video spatial grid {(height, width)} for " f"tokens_per_frame={tokens_per_frame}." ) pe_small = self._build_2d_sincos_spatial_pe( height=height, width=width, dim=projection.shape[1], device=k_video.device, dtype=k_video.dtype, ) pe_projected = torch.matmul( pe_small, projection[action_layer_idx].to(device=k_video.device, dtype=k_video.dtype), ) video_tokens = k_video.shape[1] spatial_idx = torch.arange(video_tokens, device=k_video.device) % tokens_per_frame return pe_projected.index_select(0, spatial_idx).unsqueeze(0) def _build_head_fused_kv_sin2d_pe( self, k_video: torch.Tensor, action_layer_idx: int, video_freqs: torch.Tensor, ) -> torch.Tensor: """Map fixed 2D sin/cos positions through the HeadFusedKV PE MLP.""" pos_mlps = getattr(self, "head_fused_kv_sin2d_pe_mlps", None) if pos_mlps is None: raise ValueError("HeadFusedKV+Sin2DPE requires initialized PE MLPs.") tokens_per_frame = self._infer_video_tokens_per_frame(video_freqs) height, width = self._infer_video_spatial_grid_size( video_freqs=video_freqs, tokens_per_frame=tokens_per_frame, ) if height * width != tokens_per_frame: raise ValueError( f"Invalid inferred video spatial grid {(height, width)} for " f"tokens_per_frame={tokens_per_frame}." ) pe_small = self._build_2d_sincos_spatial_pe( height=height, width=width, dim=self.new_fused_kv_pos_embed_dim, device=k_video.device, dtype=k_video.dtype, ) video_tokens = k_video.shape[1] spatial_idx = torch.arange(video_tokens, device=k_video.device) % tokens_per_frame return pos_mlps[action_layer_idx](pe_small.index_select(0, spatial_idx)).unsqueeze(0) def _apply_simple_pe_to_video_k( self, k_video: torch.Tensor, action_layer_idx: int, video_freqs: torch.Tensor, ) -> torch.Tensor: """Add spatial-only fixed 2D sin/cos PE through the learned projection.""" return k_video + self._build_simple_pe_for_video_k( k_video=k_video, action_layer_idx=action_layer_idx, video_freqs=video_freqs, ) def _allocate_aligned_3d_head_anchor_indices( self, device: torch.device, ) -> torch.Tensor: """Allocate action-attention heads to camera-view anchors.""" num_anchors = len(self.aligned_3d_action_spatial_anchors) if num_anchors == 1: return torch.zeros(self.num_heads, device=device, dtype=torch.long) if self.num_heads % 2 != 0: raise ValueError( "aligned_3d main-half head allocation requires an even " f"num_heads, got {self.num_heads}." ) num_main_heads = self.num_heads // 2 num_wrist_heads = self.num_heads - num_main_heads num_wrists = num_anchors - 1 wrist_base, wrist_remainder = divmod(num_wrist_heads, num_wrists) head_anchor_indices = [0] * num_main_heads for wrist_idx in range(num_wrists): wrist_count = wrist_base + int(wrist_idx < wrist_remainder) head_anchor_indices.extend([wrist_idx + 1] * wrist_count) return torch.tensor(head_anchor_indices, device=device, dtype=torch.long) def _build_camera_region_spatial_masks( self, height: int, width: int, device: torch.device, ) -> torch.Tensor: """Build a disjoint spatial-token mask for each camera region.""" layout = self.aligned_3d_action_spatial_anchor_layout num_cameras = len(self.aligned_3d_action_spatial_anchors) if layout == "center": return torch.ones(1, height * width, dtype=torch.bool, device=device) if layout in {"horizontal", "libero", "libero_wrist_grounded"}: if width < 2: raise ValueError( f"aligned_3dp layout {layout!r} requires width >= 2, " f"got width={width}." ) mid_w = width // 2 regions = ( (0, height, 0, mid_w), (0, height, mid_w, width), ) elif layout == "robotwin": if height < 2 or width < 2: raise ValueError( "aligned_3dp layout 'robotwin' requires height >= 2 and " f"width >= 2, got {(height, width)}." ) main_h = (height * 2) // 3 mid_w = width // 2 regions = ( (0, main_h, 0, width), (main_h, height, 0, mid_w), (main_h, height, mid_w, width), ) else: raise ValueError( "aligned_3dp has no camera-region partition for " f"aligned_3d_action_spatial_anchor_layout={layout!r}." ) if len(regions) != num_cameras: raise ValueError( f"Camera-region count {len(regions)} does not match anchor count " f"{num_cameras} for layout {layout!r}." ) masks = [] for h0, h1, w0, w1 in regions: mask = torch.zeros(height, width, dtype=torch.bool, device=device) mask[h0:h1, w0:w1] = True flat = mask.reshape(-1) if not bool(flat.any()): raise ValueError( f"aligned_3dp layout {layout!r} produced an empty camera " f"region for grid {(height, width)}." ) masks.append(flat) stacked = torch.stack(masks, dim=0) if int(stacked.sum()) != height * width or not bool( (stacked.sum(dim=0) == 1).all() ): raise ValueError( f"aligned_3dp layout {layout!r} camera regions must partition " f"the frame without gaps or overlap; got coverage " f"{int(stacked.sum())} for {height * width} tokens." ) return stacked def _build_aligned_3dp_per_head_action_mask( self, base_action_mask: torch.Tensor, video_freqs: torch.Tensor, ) -> torch.Tensor: """Restrict each action-attention head to its own camera region.""" if base_action_mask.ndim != 2: raise ValueError( "aligned_3dp expects a 2D base action mask " f"[Sa, Sv + Sa], got shape {tuple(base_action_mask.shape)}" ) action_seq_len, total_kv_len = base_action_mask.shape if total_kv_len <= action_seq_len: raise ValueError( "aligned_3dp base action mask has no video columns: " f"shape {tuple(base_action_mask.shape)}" ) video_seq_len = total_kv_len - action_seq_len if video_freqs.shape[0] != video_seq_len: raise ValueError( "aligned_3dp video_freqs sequence length " f"{video_freqs.shape[0]} != action-mask video columns " f"{video_seq_len}." ) tokens_per_frame = self._infer_video_tokens_per_frame(video_freqs) height, width = self._infer_video_spatial_grid_size( video_freqs=video_freqs, tokens_per_frame=tokens_per_frame, ) if height * width != tokens_per_frame: raise ValueError( f"Invalid inferred video spatial grid {(height, width)} for " f"tokens_per_frame={tokens_per_frame}." ) region_masks = self._build_camera_region_spatial_masks( height=height, width=width, device=base_action_mask.device, ) head_anchor_indices = self._allocate_aligned_3d_head_anchor_indices( device=base_action_mask.device, ) head_spatial = region_masks.index_select(0, head_anchor_indices) spatial_idx = ( torch.arange(video_seq_len, device=base_action_mask.device) % tokens_per_frame ) head_video_visible = head_spatial[:, spatial_idx] per_head_mask = base_action_mask.unsqueeze(0).expand( self.num_heads, -1, -1 ).clone() per_head_mask[:, :, :video_seq_len] &= head_video_visible.unsqueeze(1) return per_head_mask def _build_new_fused_kv_3d_rope_freqs( self, video_freqs: torch.Tensor, action_seq_len: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Build 3D RoPE freqs for new_fused_kv video K and action Q/K. Video K keeps each video token's original spatial h/w RoPE and uses the temporal RoPE of action token 0. Each action-attention head uses the full video temporal/h/w frequency basis at one camera-view anchor. Half of the heads use the main-camera anchor; the remainder are divided evenly among wrist-camera anchors. Camera centers are represented continuously in the compressed video-token coordinate system. """ t_dim, h_dim, w_dim = self._split_3d_rope_dims() expected_rope_dim = t_dim + h_dim + w_dim if video_freqs.ndim != 3 or video_freqs.shape[1] != 1: raise ValueError( f"video_freqs must be [Sv, 1, rope_dim], got {tuple(video_freqs.shape)}" ) if video_freqs.shape[-1] != expected_rope_dim: raise ValueError( f"video_freqs rope dim {video_freqs.shape[-1]} != expected " f"{expected_rope_dim} for attn_head_dim={self.attn_head_dim}" ) temporal_freqs = precompute_freqs_cis( self.attn_head_dim - 2 * (self.attn_head_dim // 3), end=max(action_seq_len, 1), ).to(device=video_freqs.device) video_anchor_temporal = temporal_freqs[:1].to(dtype=video_freqs.dtype) video_rope_freqs = video_freqs.clone() video_rope_freqs[:, :, :t_dim] = video_anchor_temporal.view(1, 1, t_dim) tokens_per_frame = self._infer_video_tokens_per_frame(video_freqs) height, width = self._infer_video_spatial_grid_size( video_freqs=video_freqs, tokens_per_frame=tokens_per_frame, ) if height * width != tokens_per_frame: raise ValueError( f"Invalid inferred video spatial grid {(height, width)} for " f"tokens_per_frame={tokens_per_frame}." ) # Recover the supplied video's spatial frequency bases at coordinate 1. # Their angles are all in [0, 1], so torch.angle does not wrap them. if height > 1: h_unit = video_freqs[width, 0, t_dim:t_dim + h_dim] else: h_unit = precompute_freqs_cis( self.attn_head_dim // 3, end=2 )[1].to(device=video_freqs.device, dtype=video_freqs.dtype) if width > 1: w_unit = video_freqs[1, 0, t_dim + h_dim:] else: w_unit = precompute_freqs_cis( self.attn_head_dim // 3, end=2 )[1].to(device=video_freqs.device, dtype=video_freqs.dtype) anchor_positions = torch.tensor( self.aligned_3d_action_spatial_anchors, device=video_freqs.device, dtype=torch.angle(h_unit).dtype, ) anchor_h_coordinates = anchor_positions[:, 0] * height - 0.5 anchor_w_coordinates = anchor_positions[:, 1] * width - 0.5 h_angles = anchor_h_coordinates[:, None] * torch.angle(h_unit)[None, :] w_angles = anchor_w_coordinates[:, None] * torch.angle(w_unit)[None, :] anchor_h_freqs = torch.polar(torch.ones_like(h_angles), h_angles) anchor_w_freqs = torch.polar(torch.ones_like(w_angles), w_angles) anchor_spatial = torch.cat([anchor_h_freqs, anchor_w_freqs], dim=-1) head_anchor_indices = self._allocate_aligned_3d_head_anchor_indices( device=video_freqs.device, ) action_spatial = anchor_spatial.index_select(0, head_anchor_indices) action_temporal = temporal_freqs[:action_seq_len].to(dtype=video_freqs.dtype) action_rope_freqs = torch.cat( [ action_temporal.view(action_seq_len, 1, t_dim).expand( -1, self.num_heads, -1 ), action_spatial.unsqueeze(0).expand(action_seq_len, -1, -1), ], dim=-1, ) return video_rope_freqs, action_rope_freqs # ------------------------------------------------------------------ # EEF-relative camera RoPE (`ee_rope` / `exclusive_ee_rope`). # # Rule B: a visual token's coordinate is its own position minus the anchor # of THAT TOKEN'S OWN CAMERA -- never the attending head's camera. No head # "carries" an anchor, so neither frequency tensor has a head axis: # # visual [B, Sv, 1, rope_dim] (batch: anchors differ per sample) # action [Sa, 1, rope_dim] (every action token sits at the origin) # # The two modes share every line below. They diverge at exactly one point: # `_build_exclusive_ee_rope_action_mask`, which only `exclusive_ee_rope` # calls. That single divergence is what makes the arm-2-vs-arm-3 ablation a # controlled comparison, so it is asserted by a bitwise-equality test. # ------------------------------------------------------------------ def _ee_rope_camera_of_each_video_token( self, video_seq_len: int, tokens_per_frame: int, height: int, width: int, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Per-token camera id, row, and camera-local column. Video tokens flatten as ``(frame, row, col)`` with column changing fastest, so each camera's tokens are **interleaved, not contiguous**: for a 7x14 composite the main camera owns 0-6, 14-20, 28-34, ... A contiguous ``arange(49)`` split silently hands the main group 21 wrist tokens and drops 21 real main tokens while still training plausibly. """ if width % 2 != 0: raise ValueError( f"ee_rope requires an even composite token width, got {width}" ) local_w = width // 2 index = torch.arange(video_seq_len, device=device) within_frame = index % tokens_per_frame row = torch.div(within_frame, width, rounding_mode="floor") col = within_frame % width camera = (col >= local_w).long() local_col = col - camera * local_w return camera, row.to(torch.float64), local_col.to(torch.float64) def _build_ee_rope_freqs( self, video_freqs: torch.Tensor, action_seq_len: int, eef_anchor_token: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Build EEF-relative visual/action RoPE frequencies (Rule B). Args: video_freqs: native video freqs ``[Sv, 1, rope_dim]``. action_seq_len: number of action tokens. eef_anchor_token: ``[B, 2, 2]`` continuous camera-local anchors, ordered ``(main, wrist)`` on axis 1 and **``(y, x)`` -- row first** on axis 2. The diagnostic tooling returns ``(x, y)``; a swap is silent, so callers must transpose deliberately. Returns: ``(visual [B, Sv, 1, rope_dim], action [Sa, 1, rope_dim])``. """ t_dim, h_dim, w_dim = self._split_3d_rope_dims() expected_rope_dim = t_dim + h_dim + w_dim if video_freqs.ndim != 3 or video_freqs.shape[1] != 1: raise ValueError( f"ee_rope video_freqs must be [Sv, 1, rope_dim], got " f"{tuple(video_freqs.shape)}" ) if video_freqs.shape[-1] != expected_rope_dim: raise ValueError( f"ee_rope video_freqs rope dim {video_freqs.shape[-1]} != " f"expected {expected_rope_dim} for attn_head_dim=" f"{self.attn_head_dim}" ) if eef_anchor_token.ndim != 3 or eef_anchor_token.shape[1:] != (2, 2): raise ValueError( f"eef_anchor_token must be [B, 2, 2] (cameras=(main,wrist), " f"coords=(y,x)), got {tuple(eef_anchor_token.shape)}" ) if not torch.isfinite(eef_anchor_token).all(): raise ValueError( "eef_anchor_token contains nonfinite values; an invalid " "projection must fail at resolution time, never reach RoPE." ) device = video_freqs.device video_seq_len = video_freqs.shape[0] tokens_per_frame = self._infer_video_tokens_per_frame(video_freqs) height, width = self._infer_video_spatial_grid_size( video_freqs=video_freqs, tokens_per_frame=tokens_per_frame, ) if height * width != tokens_per_frame: raise ValueError( f"Invalid inferred video spatial grid {(height, width)} for " f"tokens_per_frame={tokens_per_frame}." ) # Recover the per-dimension angular bases exactly as the aligned_3d # builder does, so both share one definition of "one token of travel". if height > 1: h_unit = video_freqs[width, 0, t_dim:t_dim + h_dim] else: h_unit = precompute_freqs_cis( self.attn_head_dim // 3, end=2 )[1].to(device=device, dtype=video_freqs.dtype) if width > 1: w_unit = video_freqs[1, 0, t_dim + h_dim:] else: w_unit = precompute_freqs_cis( self.attn_head_dim // 3, end=2 )[1].to(device=device, dtype=video_freqs.dtype) h_angle_unit = torch.angle(h_unit) w_angle_unit = torch.angle(w_unit) camera, token_row, token_local_col = self._ee_rope_camera_of_each_video_token( video_seq_len=video_seq_len, tokens_per_frame=tokens_per_frame, height=height, width=width, device=device, ) anchors = eef_anchor_token.to(device=device, dtype=h_angle_unit.dtype) # Gather each token's OWN camera anchor -- this is Rule B. anchor_row = anchors[:, :, 0].index_select(1, camera) # [B, Sv] anchor_col = anchors[:, :, 1].index_select(1, camera) # [B, Sv] rel_row = token_row.unsqueeze(0).to(anchor_row.dtype) - anchor_row rel_col = token_local_col.unsqueeze(0).to(anchor_col.dtype) - anchor_col h_angles = rel_row.unsqueeze(-1) * h_angle_unit.view(1, 1, h_dim) w_angles = rel_col.unsqueeze(-1) * w_angle_unit.view(1, 1, w_dim) visual_h = torch.polar(torch.ones_like(h_angles), h_angles) visual_w = torch.polar(torch.ones_like(w_angles), w_angles) temporal_freqs = precompute_freqs_cis( self.attn_head_dim - 2 * (self.attn_head_dim // 3), end=max(action_seq_len, 1), ).to(device=device) # Visual K is read by action queries, so it sits at action time 0, # matching the aligned_3d convention. visual_temporal = temporal_freqs[:1].to(dtype=video_freqs.dtype) visual_temporal = visual_temporal.view(1, 1, t_dim).expand( anchors.shape[0], video_seq_len, t_dim ) visual_rope_freqs = torch.cat( [ visual_temporal, visual_h.to(video_freqs.dtype), visual_w.to(video_freqs.dtype), ], dim=-1, ).unsqueeze(2) # [B, Sv, 1, D] # Every action token sits at the spatial ORIGIN: the anchor is the # origin, so its relative coordinate is (0, 0) and the spatial phase is # identically 1. This is why action freqs carry no head axis. action_temporal = temporal_freqs[:action_seq_len].to(dtype=video_freqs.dtype) action_spatial = torch.ones( action_seq_len, h_dim + w_dim, device=device, dtype=video_freqs.dtype, ) action_rope_freqs = torch.cat( [action_temporal.view(action_seq_len, t_dim), action_spatial], dim=-1, ).unsqueeze(1) # [Sa, 1, D] return visual_rope_freqs, action_rope_freqs def _build_exclusive_ee_rope_action_mask( self, base_action_mask: torch.Tensor, video_freqs: torch.Tensor, video_seq_len: int, action_seq_len: int, ) -> torch.Tensor: """Restrict each action head to one camera's first-frame visual tokens. Heads ``0..H/2-1`` see main-camera tokens, ``H/2..H-1`` see wrist. All heads keep every action key. Later visual frames stay masked because the incoming 2D mask already excludes them; this only ever removes visibility, never adds it. Returns ``[num_heads, Sa, Sv + Sa]``, which broadcasts against SDPA's ``[B, num_heads, Sq, Sk]`` without touching the attention kernel. """ if base_action_mask.ndim != 2: raise ValueError( "exclusive_ee_rope expects a 2D base action mask to refine, " f"got ndim={base_action_mask.ndim}" ) expected = (action_seq_len, video_seq_len + action_seq_len) if base_action_mask.shape != torch.Size(expected): raise ValueError( f"exclusive_ee_rope base action mask must be {expected}, got " f"{tuple(base_action_mask.shape)}" ) if self.num_heads % 2 != 0: raise ValueError( "exclusive_ee_rope requires an even num_heads for the 50/50 " f"camera split, got {self.num_heads}" ) device = base_action_mask.device tokens_per_frame = self._infer_video_tokens_per_frame(video_freqs) height, width = self._infer_video_spatial_grid_size( video_freqs=video_freqs, tokens_per_frame=tokens_per_frame, ) camera, _, _ = self._ee_rope_camera_of_each_video_token( video_seq_len=video_seq_len, tokens_per_frame=tokens_per_frame, height=height, width=width, device=device, ) num_main_heads = self.num_heads // 2 head_camera = torch.cat( [ torch.zeros(num_main_heads, device=device, dtype=torch.long), torch.ones( self.num_heads - num_main_heads, device=device, dtype=torch.long ), ] ) # [num_heads, Sv] -- True where the token belongs to that head's camera. visual_visible = head_camera[:, None] == camera[None, :] action_visible = torch.ones( self.num_heads, action_seq_len, device=device, dtype=torch.bool ) head_key_visible = torch.cat([visual_visible, action_visible], dim=1) return base_action_mask.unsqueeze(0) & head_key_visible[:, None, :] # ------------------------------------------------------------------ # Shared building blocks (2026-07-05 dedup). The validation preamble, # the raw-KV video loop, and the action loop previously existed as # 2/4/4 near-identical copies across forward_decoupled, # _forward_decoupled_fused_rope, prefill_video_kv, and # forward_action_with_video_kv; a fix landing in one copy silently # missed the others. The four public methods now differ ONLY in how # the per-action-layer video K/V is sourced (selected/mixed stack vs # kv_fusion output). # ------------------------------------------------------------------ def _validate_decoupled_forward_inputs( self, embeds_all: Dict[str, torch.Tensor], attention_masks: Dict[str, torch.Tensor], freqs_all: Dict[str, torch.Tensor], t_mod_all: Dict[str, torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: """Validate per-expert forward inputs; return (video_mask, action_mask).""" for key in (VIDEO_EXPERT_KEY, ACTION_EXPERT_KEY): if key not in embeds_all: raise ValueError(f"Missing '{key}' in embeds_all") if key not in freqs_all: raise ValueError(f"Missing '{key}' in freqs_all") if key not in t_mod_all: raise ValueError(f"Missing '{key}' in t_mod_all") if key not in attention_masks: raise ValueError(f"Missing '{key}' in attention_masks") video_mask = attention_masks[VIDEO_EXPERT_KEY] action_mask = attention_masks[ACTION_EXPERT_KEY] if video_mask.ndim != 2 or video_mask.shape[0] != video_mask.shape[1]: raise ValueError( f"Video attention mask must be square 2D, got shape " f"{tuple(video_mask.shape)}" ) video_seq_len = embeds_all[VIDEO_EXPERT_KEY].shape[1] action_seq_len = embeds_all[ACTION_EXPERT_KEY].shape[1] if video_mask.shape[0] != video_seq_len: raise ValueError( f"Video mask seq length {video_mask.shape[0]} != " f"video token seq length {video_seq_len}" ) expected_action_mask_shape = (action_seq_len, video_seq_len + action_seq_len) if action_mask.shape != torch.Size(expected_action_mask_shape): raise ValueError( f"Action attention mask must have shape {expected_action_mask_shape}, " f"got {tuple(action_mask.shape)}" ) return video_mask, action_mask def _run_video_layers_raw_kv( self, x: torch.Tensor, video_freqs: torch.Tensor, video_t_mod: torch.Tensor, video_context: Optional[dict], video_mask: torch.Tensor, ) -> tuple[torch.Tensor, list[torch.Tensor], list[torch.Tensor]]: """Run every video layer (3D-RoPE'd self-attn), caching video K and V. By default the cached K is pre-RoPE (raw / pre-norm, depending on the projection mode). ``legacy_3d`` instead caches the post-3D-RoPE K used in video self-attention, so fusion sees already-positioned keys. The full post-block MUST run each layer because layer i+1's K depends on layer i's post-block output. Returns: (final video tokens, per-layer K list, per-layer V list); cached tensors are flattened ``[B, Sv, H*Dh]``. """ video_expert = self.mixtures[VIDEO_EXPERT_KEY] raw_k_per_layer: list[torch.Tensor] = [] raw_v_per_layer: list[torch.Tensor] = [] cache_post_rope_k = self._cache_post_3d_rope_video_k for video_layer_idx in range(self.video_num_layers): video_block = video_expert.blocks[video_layer_idx] build_attention_io = ( self._build_expert_attention_io_with_raw_k if ( self.kv_source_mode != "new_fused_kv" or self.new_fused_kv_projection_mode in { "HeadFusedKV", "HeadFusedKV+Sin2DPE", "MLPMixerFusedKV", } ) else self._build_expert_attention_io_with_pre_norm_k ) ( q_video, k_video_rope, k_video_raw, v_video, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gradient_checkpointing, ) = build_attention_io( expert=video_expert, block=video_block, x=x, freqs=video_freqs, t_mod=video_t_mod, ) mixed = self._mixed_attention( q_cat=q_video, k_cat=k_video_rope, v_cat=v_video, attention_mask=video_mask, ) x = self._apply_post_with_optional_checkpoint( block=video_block, residual_x=residual_x, gate_msa=gate_msa, shift_mlp=shift_mlp, scale_mlp=scale_mlp, gate_mlp=gate_mlp, use_gradient_checkpointing=use_gradient_checkpointing, mixed_slice=mixed, context_payload=video_context, ) raw_k_per_layer.append( k_video_rope if cache_post_rope_k else k_video_raw ) raw_v_per_layer.append(v_video) return x, raw_k_per_layer, raw_v_per_layer def _validate_cached_action_mask( self, *, action_mask: torch.Tensor, video_seq_len: int, action_seq_len: int, use_new_fused_kv: bool, ) -> None: """Check the action mask against the sourced video K on the cache path. Only the KV-cache inference path calls this: training validates its square joint mask upfront. Subclasses whose mask is request-dependent override this to admit their own batched form. """ expected_mask_shape = (action_seq_len, video_seq_len + action_seq_len) # exclusive_ee_rope and aligned_3dp refine the 2D mask into a # head-aware [num_heads, Sa, Sk] mask above. head_aware_ok = ( use_new_fused_kv and self.new_fused_kv_rope_mode in {"exclusive_ee_rope", "aligned_3dp"} ) if head_aware_ok: expected_head_shape = (self.num_heads, *expected_mask_shape) if action_mask.ndim != 3: raise ValueError( f"{self.new_fused_kv_rope_mode} " "action_attention_mask must be 3D " f"{expected_head_shape}, got ndim={action_mask.ndim}" ) if action_mask.shape != torch.Size(expected_head_shape): raise ValueError( f"{self.new_fused_kv_rope_mode} " "action_attention_mask shape must be " f"{expected_head_shape}, got {tuple(action_mask.shape)}" ) else: if action_mask.ndim != 2: raise ValueError( f"action_attention_mask must be 2D, " f"got ndim={action_mask.ndim}" ) if action_mask.shape != torch.Size(expected_mask_shape): raise ValueError( f"action_attention_mask shape must be {expected_mask_shape}, " f"got {tuple(action_mask.shape)}" ) def _run_action_layers_with_video_kv_source( self, x: torch.Tensor, action_freqs: torch.Tensor, action_t_mod: torch.Tensor, action_context: Optional[dict], action_mask: torch.Tensor, video_kv_for_layer: Callable[[int], tuple[torch.Tensor, torch.Tensor]], validate_mask_at_layer0: bool = False, video_freqs: Optional[torch.Tensor] = None, eef_anchor_token: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Run every action layer against per-layer sourced raw video K/V. ``video_kv_for_layer(action_layer_idx)`` returns the cached video K and the video V for that action layer. For most modes that K is RAW (un-RoPE'd); for ``legacy_3d`` it is already post-3D-RoPE. For ``new_fused_kv``, ``aligned_3d`` / ``aligned_3dp`` apply video-style 3D RoPE to both video K and action Q/K (``aligned_3dp`` additionally restricts each head to its camera-region video tokens), ``aligned_1d`` applies temporal-only 1D RoPE to both fused video K and action Q/K, ``original_3d`` applies the original video 3D RoPE to fused video K and keeps action 1D RoPE, while ``video_zero_1d`` / ``legacy_3d`` apply no extra video RoPE in this action path. V is never RoPE'd. Args: validate_mask_at_layer0: When True (KV-cache inference path, where no square video mask was validated upfront), check the action mask shape against the sourced video K at the first layer. """ action_expert = self.mixtures[ACTION_EXPERT_KEY] use_new_fused_kv = self.kv_source_mode == "new_fused_kv" use_aligned_new_fused_kv_rope = ( use_new_fused_kv and self.new_fused_kv_rope_mode in self._aligned_3d_family_rope_modes ) use_aligned_3dp_camera_mask = ( use_new_fused_kv and self.new_fused_kv_rope_mode == "aligned_3dp" ) use_aligned_1d_new_fused_kv_rope = ( use_new_fused_kv and self.new_fused_kv_rope_mode == "aligned_1d" ) use_original_new_fused_kv_rope = ( use_new_fused_kv and self.new_fused_kv_rope_mode == "original_3d" ) # video_zero_1d: fuse pre-RoPE K, skip action-path video RoPE. # legacy_3d: fuse post-3D-RoPE K, skip action-path video RoPE (same skip). use_skip_action_path_video_rope = ( use_new_fused_kv and self.new_fused_kv_rope_mode in {"video_zero_1d", "legacy_3d"} ) use_ee_rope = ( use_new_fused_kv and self.new_fused_kv_rope_mode in self.ee_rope_modes ) if use_ee_rope: if video_freqs is None: raise ValueError( f"new_fused_kv {self.new_fused_kv_rope_mode!r} action path " "requires video_freqs to build EEF-relative RoPE." ) if eef_anchor_token is None: raise ValueError( f"new_fused_kv {self.new_fused_kv_rope_mode!r} requires " "eef_anchor_token [B, 2, 2]; it must be threaded from " "build_inputs() and never defaulted." ) video_rope_freqs, action_rope_freqs = self._build_ee_rope_freqs( video_freqs=video_freqs, action_seq_len=x.shape[1], eef_anchor_token=eef_anchor_token, ) # THE single divergence point between the two sibling modes. if self.new_fused_kv_rope_mode == "exclusive_ee_rope": action_mask = self._build_exclusive_ee_rope_action_mask( base_action_mask=action_mask, video_freqs=video_freqs, video_seq_len=video_freqs.shape[0], action_seq_len=x.shape[1], ) elif use_aligned_new_fused_kv_rope: if video_freqs is None: raise ValueError( "new_fused_kv fixed-RoPE action path requires video_freqs " "to build video-style 3D RoPE for action attention." ) video_rope_freqs, action_rope_freqs = self._build_new_fused_kv_3d_rope_freqs( video_freqs=video_freqs, action_seq_len=x.shape[1], ) if use_aligned_3dp_camera_mask: action_mask = self._build_aligned_3dp_per_head_action_mask( base_action_mask=action_mask, video_freqs=video_freqs, ) elif use_original_new_fused_kv_rope: if video_freqs is None: raise ValueError( "new_fused_kv original-3D-RoPE action path requires video_freqs " "to apply the original video 3D RoPE to fused video K." ) # Rank/rope-dim validation for the original_3d branch. Unlike the # aligned_3d path, original_3d uses video_freqs verbatim as the RoPE # basis for rope_apply(k_video_raw, video_rope_freqs, num_heads), so a # malformed freqs tensor (wrong rank, or a last dim that does not match # attn_head_dim // 2) would silently mis-rotate every video token. This # mirrors the shape contract asserted by the aligned-path builder # _build_new_fused_kv_3d_rope_freqs (t_dim + h_dim + w_dim). It is # complementary to the 5fb5881 length guard (freqs.shape[0] vs Sv), # which is enforced separately at the per-layer rope_apply site. t_dim, h_dim, w_dim = self._split_3d_rope_dims() expected_rope_dim = t_dim + h_dim + w_dim if video_freqs.ndim != 3 or video_freqs.shape[1] != 1: raise ValueError( "new_fused_kv original_3d video_freqs must be " "[Sv, 1, rope_dim], got " f"{tuple(video_freqs.shape)}" ) if video_freqs.shape[-1] != expected_rope_dim: raise ValueError( "new_fused_kv original_3d video_freqs rope dim " f"{video_freqs.shape[-1]} != expected {expected_rope_dim} " f"for attn_head_dim={self.attn_head_dim}" ) video_rope_freqs = video_freqs action_rope_freqs = action_freqs elif use_aligned_1d_new_fused_kv_rope: video_rope_freqs = action_freqs[:1] action_rope_freqs = action_freqs elif use_skip_action_path_video_rope: video_rope_freqs = None action_rope_freqs = action_freqs else: video_rope_freqs = action_freqs[:1] action_rope_freqs = action_freqs for action_layer_idx in range(self.action_num_layers): k_video_raw, v_video = video_kv_for_layer(action_layer_idx) if ( use_new_fused_kv and self.new_fused_kv_projection_mode in {"simple+PE", "simple+PE-postnorm"} ): if video_freqs is None: raise ValueError( f"{self.new_fused_kv_projection_mode} new_fused_kv " "requires video_freqs to build spatial 2D sin/cos PE." ) if self.new_fused_kv_projection_mode == "simple+PE": k_video_raw = self._apply_simple_pe_to_video_k( k_video=k_video_raw, action_layer_idx=action_layer_idx, video_freqs=video_freqs, ) if self.k_fused_norm is not None: k_video_raw = self.k_fused_norm[action_layer_idx](k_video_raw) else: k_video_raw = k_video_raw + self._build_simple_pe_for_video_k( k_video=k_video_raw, action_layer_idx=action_layer_idx, video_freqs=video_freqs, ) if ( use_new_fused_kv and self.new_fused_kv_projection_mode == "HeadFusedKV+Sin2DPE" ): if video_freqs is None: raise ValueError( "HeadFusedKV+Sin2DPE requires video_freqs to build spatial 2D sin/cos PE." ) k_video_raw = k_video_raw + self._build_head_fused_kv_sin2d_pe( k_video=k_video_raw, action_layer_idx=action_layer_idx, video_freqs=video_freqs, ) if self.k_fused_norm is not None: k_video_raw = self.k_fused_norm[action_layer_idx](k_video_raw) if use_new_fused_kv: if use_skip_action_path_video_rope: k_video = k_video_raw else: assert video_rope_freqs is not None k_video = rope_apply(k_video_raw, video_rope_freqs, self.num_heads) else: k_video = self._apply_action_zero_rope_to_video_k( k_video_raw=k_video_raw, action_freqs=video_rope_freqs, ) if validate_mask_at_layer0 and action_layer_idx == 0: self._validate_cached_action_mask( action_mask=action_mask, video_seq_len=k_video.shape[1], action_seq_len=x.shape[1], use_new_fused_kv=use_new_fused_kv, ) action_block = action_expert.blocks[action_layer_idx] ( q_action, k_action, v_action, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gradient_checkpointing, ) = self._build_expert_attention_io( expert=action_expert, block=action_block, x=x, freqs=action_rope_freqs, t_mod=action_t_mod, ) k_cat = torch.cat([k_video, k_action], dim=1) v_cat = torch.cat([v_video, v_action], dim=1) mixed = self._mixed_attention( q_cat=q_action, k_cat=k_cat, v_cat=v_cat, attention_mask=action_mask, ) x = self._apply_post_with_optional_checkpoint( block=action_block, residual_x=residual_x, gate_msa=gate_msa, shift_mlp=shift_mlp, scale_mlp=scale_mlp, gate_mlp=gate_mlp, use_gradient_checkpointing=use_gradient_checkpointing, mixed_slice=mixed, context_payload=action_context, ) return x def forward_decoupled( self, embeds_all: Dict[str, torch.Tensor], attention_masks: Dict[str, torch.Tensor], freqs_all: Dict[str, torch.Tensor], context_all: Dict[str, Optional[dict]], t_mod_all: Dict[str, torch.Tensor], eef_anchor_token: Optional[torch.Tensor] = None, ) -> Dict[str, torch.Tensor]: """Training forward with raw video-K fusion and action-aligned RoPE. ``eef_anchor_token`` is ``[B, 2, 2]`` and required only by the ``ee_rope`` / ``exclusive_ee_rope`` modes; every other mode ignores it. """ # Fused-MLP mode routes through a dedicated fuse-then-RoPE path. Selected # modes (final_only / uniform_end / fused_kv) keep the body below. if self.kv_fusion is not None: return self._forward_decoupled_fused_rope( embeds_all=embeds_all, attention_masks=attention_masks, freqs_all=freqs_all, context_all=context_all, t_mod_all=t_mod_all, eef_anchor_token=eef_anchor_token, ) video_mask, action_mask = self._validate_decoupled_forward_inputs( embeds_all=embeds_all, attention_masks=attention_masks, freqs_all=freqs_all, t_mod_all=t_mod_all, ) x_video, raw_k_per_layer, raw_v_per_layer = self._run_video_layers_raw_kv( x=embeds_all[VIDEO_EXPERT_KEY], video_freqs=freqs_all[VIDEO_EXPERT_KEY], video_t_mod=t_mod_all[VIDEO_EXPERT_KEY], video_context=context_all.get(VIDEO_EXPERT_KEY), video_mask=video_mask, ) video_kv_cache = [ {"k": k, "v": v} for k, v in zip(raw_k_per_layer, raw_v_per_layer) ] stacked_k, stacked_v = self._stack_video_kv(video_kv_cache) x_action = self._run_action_layers_with_video_kv_source( x=embeds_all[ACTION_EXPERT_KEY], action_freqs=freqs_all[ACTION_EXPERT_KEY], action_t_mod=t_mod_all[ACTION_EXPERT_KEY], action_context=context_all.get(ACTION_EXPERT_KEY), action_mask=action_mask, video_kv_for_layer=lambda idx: self._select_or_mix_stacked_video_kv( stacked_k=stacked_k, stacked_v=stacked_v, action_layer_idx=idx, ), video_freqs=freqs_all[VIDEO_EXPERT_KEY], eef_anchor_token=eef_anchor_token, ) return { VIDEO_EXPERT_KEY: x_video, ACTION_EXPERT_KEY: x_action, } def _forward_decoupled_fused_rope( self, embeds_all: Dict[str, torch.Tensor], attention_masks: Dict[str, torch.Tensor], freqs_all: Dict[str, torch.Tensor], context_all: Dict[str, Optional[dict]], t_mod_all: Dict[str, torch.Tensor], eef_anchor_token: Optional[torch.Tensor] = None, ) -> Dict[str, torch.Tensor]: """Training forward for ``fused_mlp`` mode with fuse-then-RoPE. Video self-attention is identical to the selected-mode path: it uses the video's native 3D RoPE. The raw (un-RoPE'd) video K and V of every video layer are cached, stacked into ``[B, Sv, N, D]``, and fused by ``self.kv_fusion`` into one K/V per action layer. Each fused, still-raw K is then re-RoPE'd right before it is concatenated with the action K (fuse-then-RoPE, per action layer). This mirrors the ``_select_or_mix_stacked_video_kv`` call in the selected-mode path, but swaps selection/mixing for the learned MLP fusion. Args: embeds_all: Per-expert input tokens (``video`` / ``action``). attention_masks: Per-expert attention masks (video square, action rectangular ``[Sa, Sv + Sa]``). freqs_all: Per-expert RoPE frequencies. context_all: Per-expert optional cross-attention context payloads. t_mod_all: Per-expert AdaLN time-modulation tensors. Returns: Dict with denoised ``video`` and ``action`` tokens. """ video_mask, action_mask = self._validate_decoupled_forward_inputs( embeds_all=embeds_all, attention_masks=attention_masks, freqs_all=freqs_all, t_mod_all=t_mod_all, ) action_freqs = freqs_all[ACTION_EXPERT_KEY] # ---- Video loop: run self-attn (3D RoPE), cache RAW K + V per layer. ---- x_video, raw_k_per_layer, raw_v_per_layer = self._run_video_layers_raw_kv( x=embeds_all[VIDEO_EXPERT_KEY], video_freqs=freqs_all[VIDEO_EXPERT_KEY], video_t_mod=t_mod_all[VIDEO_EXPERT_KEY], video_context=context_all.get(VIDEO_EXPERT_KEY), video_mask=video_mask, ) # Stack raw K/V over video layers -> [B, Sv, N, D] for the fusion module. all_k = torch.stack(raw_k_per_layer, dim=2) all_v = torch.stack(raw_v_per_layer, dim=2) fused_kv = self.kv_fusion(all_k, all_v) # list length M of {"k","v"} [B,Sv,D] # ---- Action loop: fuse-then-RoPE per action layer (the RoPE rotation # of the fused raw K happens inside the shared action-loop helper). ---- x_action = self._run_action_layers_with_video_kv_source( x=embeds_all[ACTION_EXPERT_KEY], action_freqs=action_freqs, action_t_mod=t_mod_all[ACTION_EXPERT_KEY], action_context=context_all.get(ACTION_EXPERT_KEY), action_mask=action_mask, video_kv_for_layer=lambda idx: ( fused_kv[idx]["k"], fused_kv[idx]["v"], ), video_freqs=freqs_all[VIDEO_EXPERT_KEY], eef_anchor_token=eef_anchor_token, ) return { VIDEO_EXPERT_KEY: x_video, ACTION_EXPERT_KEY: x_action, } def prefill_video_kv( self, video_tokens: torch.Tensor, video_freqs: torch.Tensor, video_t_mod: torch.Tensor, video_context_payload: Optional[dict], video_attention_mask: torch.Tensor, ) -> list[dict[str, torch.Tensor]]: """Run video layers and cache video K plus normal V for action. The return contract is mode-dependent, mirroring the two branches below: - ``kv_fusion is not None`` (fused_mlp): the K/V of every video layer are stacked into ``[B, Sv, N, D]`` and fused by ``self.kv_fusion`` into one K/V per action layer. Returns the length ``action_num_layers`` (M) fused list of ``{"k","v"}`` tensors ``[B, Sv, D]``, exactly what ``forward_action_with_video_kv`` consumes in fused mode. - ``kv_fusion is None`` (selected modes): returns the length ``video_num_layers`` (N) list of per-layer ``{"k","v"}`` caches, to be stacked/selected downstream by ``forward_action_with_video_kv``. By default the cached K is RAW (un-RoPE'd) and action attention applies mode-specific RoPE later in ``forward_action_with_video_kv``. ``legacy_3d`` instead caches post-3D-RoPE K and skips extra video RoPE in the action path. V is never RoPE'd. Args: video_tokens: Video expert input tokens ``[B, Sv, D_video]``. video_freqs: Video 3D RoPE frequencies for the self-attention loop. video_t_mod: Video AdaLN time-modulation tensor. video_context_payload: Optional video cross-attention context payload. video_attention_mask: Square video self-attention mask ``[Sv, Sv]``. Returns: list of ``{"k","v"}`` dicts. Length is ``action_num_layers`` (M) when ``kv_fusion is not None`` (fused list), else ``video_num_layers`` (N) (raw per-layer list). Each tensor has shape ``[B, Sv, D]``. """ if video_attention_mask.ndim != 2: raise ValueError( f"video_attention_mask must be 2D [Sv, Sv], got shape " f"{tuple(video_attention_mask.shape)}" ) if video_attention_mask.shape[0] != video_attention_mask.shape[1]: raise ValueError( f"video_attention_mask must be square, got shape " f"{tuple(video_attention_mask.shape)}" ) if video_attention_mask.shape[0] != video_tokens.shape[1]: raise ValueError( f"video_attention_mask seq length {video_attention_mask.shape[0]} " f"!= video_tokens seq length {video_tokens.shape[1]}" ) # Run the full video block per layer (self-attn uses 3D-RoPE'd K), # caching K/V of every layer via the shared training helper. Default # modes cache RAW K; ``legacy_3d`` caches post-3D-RoPE K. _, raw_k_per_layer, raw_v_per_layer = self._run_video_layers_raw_kv( x=video_tokens, video_freqs=video_freqs, video_t_mod=video_t_mod, video_context=video_context_payload, video_mask=video_attention_mask, ) if self.kv_fusion is not None: # Stack raw K/V over video layers -> [B, Sv, N, D] for the fusion # module. dtype invariant: fused K stays in the model dtype; the # action-token-0 RoPE applied later in forward returns .to(x.dtype) # (wan_video_dit.py rope_apply), so no cast is needed here. all_k = torch.stack(raw_k_per_layer, dim=2) all_v = torch.stack(raw_v_per_layer, dim=2) # Length == action_num_layers list of {"k","v"} tensors [B, Sv, D]. return self.kv_fusion(all_k, all_v) return [ {"k": k, "v": v, "freqs": video_freqs} if self.kv_source_mode == "new_fused_kv" else {"k": k, "v": v} for k, v in zip(raw_k_per_layer, raw_v_per_layer) ] def forward_action_with_video_kv( self, video_kv_per_layer: list[dict[str, torch.Tensor]], action_tokens: torch.Tensor, action_freqs: torch.Tensor, action_t_mod: torch.Tensor, action_context_payload: Optional[dict], action_attention_mask: torch.Tensor, eef_anchor_token: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Run action using raw cached video K rotated by action token-0 RoPE. The expected length of ``video_kv_per_layer`` is mode-dependent and must match what ``prefill_video_kv`` produced: - ``kv_fusion is not None`` (fused_mlp): ``video_kv_per_layer`` is the length ``action_num_layers`` (M) FUSED list from ``prefill_video_kv`` (already one raw K/V per action layer). ``_stack_video_kv`` is NOT called (it asserts len == ``video_num_layers``); each action layer re-RoPEs its own fused raw K and uses V raw. - ``kv_fusion is None`` (selected modes): ``video_kv_per_layer`` is the length ``video_num_layers`` (N) raw per-layer list; it is stacked and the selected/mixed K/V per action layer is re-RoPE'd. Mirrors the training action loop in ``_forward_decoupled_fused_rope`` (fused mode) / ``forward_decoupled`` (selected mode). Args: video_kv_per_layer: Cached video K/V from ``prefill_video_kv``. Length ``action_num_layers`` (M) when ``kv_fusion is not None`` (fused list), else ``video_num_layers`` (N) raw per-layer list. Each ``{"k","v"}`` tensor has shape ``[B, Sv, D]``. action_tokens: Action expert input tokens ``[B, Sa, D_action]``. action_freqs: Action RoPE frequencies for the legacy fixed-RoPE path. ``new_fused_kv`` uses action length plus cached ``video_freqs`` to build video-style 3D RoPE instead. action_t_mod: Action AdaLN time-modulation tensor. action_context_payload: Optional action cross-attention context. action_attention_mask: Rectangular action mask ``[Sa, Sv + Sa]``. Returns: Denoised action tokens ``[B, Sa, D_action]``. """ if self.kv_fusion is not None: # Fused-MLP mode: video_kv_per_layer is the length-M fused list from # prefill_video_kv (already one raw K/V per action layer). Do NOT # call _stack_video_kv (it asserts len == video_num_layers). Per # action layer, the shared helper re-RoPEs the fused raw K with # mode-specific fixed RoPE and uses V raw. Mirrors the training # action loop in _forward_decoupled_fused_rope. if len(video_kv_per_layer) != self.action_num_layers: raise ValueError( f"video_kv_per_layer length {len(video_kv_per_layer)} != " f"action_num_layers {self.action_num_layers}" ) return self._run_action_layers_with_video_kv_source( x=action_tokens, action_freqs=action_freqs, action_t_mod=action_t_mod, action_context=action_context_payload, action_mask=action_attention_mask, video_kv_for_layer=lambda idx: ( video_kv_per_layer[idx]["k"], video_kv_per_layer[idx]["v"], ), validate_mask_at_layer0=True, eef_anchor_token=eef_anchor_token, ) stacked_k, stacked_v = self._stack_video_kv(video_kv_per_layer) video_freqs = None if self.kv_source_mode == "new_fused_kv": video_freqs = video_kv_per_layer[0].get("freqs") if video_freqs is None: raise ValueError( "new_fused_kv fixed-RoPE inference cache must include " "`freqs`; use prefill_video_kv() from the same model." ) # Sequence-axis derivation: the cached sidecar is consumed by # rope_apply(k_video_raw, video_rope_freqs, num_heads) at # _run_action_layers_with_video_kv_source line ~414. There # k_video_raw is [B, Sv, D] -> rearranged to [B, Sv, n, d] and the # freqs (video_rope_freqs = video_freqs.clone(), so same axis-0 # length) broadcast as [Sv, 1, rope_dim/2] against [B, Sv, n, d/2]. # rope_apply broadcasts positionally, so a shorter/longer freqs # axis-0 would either error or (when equal by coincidence) silently # mis-rotate; axis 0 of video_freqs MUST equal the cached video K # sequence length Sv. _build_new_fused_kv_3d_rope_freqs asserts # video_freqs is [Sv, 1, rope_dim] (line ~217), fixing Sv on axis 0. # stacked_k is [N, B, Sv, num_heads, attn_head_dim] from # _stack_video_kv, so Sv = stacked_k.shape[2]. cached_video_seq_len = stacked_k.shape[2] freqs_seq_len = video_freqs.shape[0] if freqs_seq_len != cached_video_seq_len: raise ValueError( "new_fused_kv cached `freqs` sequence length " f"{freqs_seq_len} != cached video K sequence length " f"{cached_video_seq_len}; the freqs sidecar is malformed or " "stale. Regenerate the cache with prefill_video_kv() from " "the same model." ) return self._run_action_layers_with_video_kv_source( x=action_tokens, action_freqs=action_freqs, action_t_mod=action_t_mod, action_context=action_context_payload, action_mask=action_attention_mask, video_kv_for_layer=lambda idx: self._select_or_mix_stacked_video_kv( stacked_k=stacked_k, stacked_v=stacked_v, action_layer_idx=idx, ), validate_mask_at_layer0=True, video_freqs=video_freqs, eef_anchor_token=eef_anchor_token, ) class FasterWAMDecoupled(FastWAMDecoupled): """FastWAMDecoupled variant using action-aligned video K RoPE.""" def build_inputs(self, sample, tiled: bool = False): """Extend the base inputs with the EEF anchor when a mode needs it. Overridden here rather than in ``FastWAM.build_inputs`` (upstream) or in ``FastWAMDecoupled`` (which never runs these modes): the EEF-relative RoPE modes are hosted by this class, so the input contract belongs with them. Adds one key and changes nothing else. ``eef_anchor_token`` is ``[B, 2, 2]`` float32, cameras ordered ``(main, wrist)`` and coordinates ``(y, x)`` -- row first. The ``eef_anchor_observed`` flag is deliberately NOT threaded here: nothing in attention, masking, or frequency construction reads it, and an unread tensor in the model contract is exactly the field a later reader re-interprets as a rejection gate (plan Sections 16.2, 19.2). """ inputs = super().build_inputs(sample, tiled=tiled) anchor = sample.get("eef_anchor_token") if anchor is None: return inputs if not torch.is_tensor(anchor): anchor = torch.as_tensor(anchor) if anchor.ndim != 3 or anchor.shape[1:] != (2, 2): raise ValueError( "`sample['eef_anchor_token']` must be [B, 2, 2] with cameras " f"(main, wrist) and coords (y, x); got {tuple(anchor.shape)}" ) if not torch.isfinite(anchor).all(): raise ValueError( "`sample['eef_anchor_token']` contains nonfinite values; an " "invalid projection must fail during load-time resolution, " "never reach the model (plan Section 19.1 Decision 1)." ) inputs["eef_anchor_token"] = anchor.to( device=self.device, dtype=torch.float32, non_blocking=True ) return inputs @classmethod def from_wan22_pretrained( cls, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16, model_id: str = "Wan-AI/Wan2.2-TI2V-5B", tokenizer_model_id: str = "Wan-AI/Wan2.1-T2V-1.3B", tokenizer_max_len: int = 512, load_text_encoder: bool = True, proprio_dim: Optional[int] = None, redirect_common_files: bool = True, video_dit_config: dict[str, Any] | None = None, action_dit_config: dict[str, Any] | None = None, action_dit_pretrained_path: str | None = None, skip_dit_load_from_pretrain: bool = False, mot_checkpoint_mixed_attn: bool = True, video_train_shift: float = 5.0, video_infer_shift: float = 5.0, video_num_train_timesteps: int = 1000, action_train_shift: float = 5.0, action_infer_shift: float = 5.0, action_num_train_timesteps: int = 1000, loss_lambda_video: float = 1.0, loss_lambda_action: float = 1.0, decoupled: bool = True, kv_source_mapping: list[int] | None = None, kv_source_mode: str = "final_only", fixed_rope: bool = True, new_fused_kv_rope_mode: str = "aligned_3d", aligned_3d_action_spatial_anchor_layout: str | None = "center", eef_calibration_path: str | None = None, eef_raw_source_resolution: int | None = None, eef_processed_video_size: list[int] | tuple[int, int] | None = None, new_fused_kv_projection_mode: str = "full", new_fused_kv_pos_embed_max_tokens: int = 4096, new_fused_kv_pos_embed_dim: int = 128, new_fused_kv_mlp_mixer_num_blocks: int = 1, new_fused_kv_mlp_mixer_token_mlp_ratio: float = 4.0, new_fused_kv_mlp_mixer_channel_mlp_ratio: float = 4.0, new_fused_kv_head_fused_kv_sin2d_pe_mlp_hidden_dim: int = 512, new_fused_kv_simple_head_softmax_fuse_mode: str = "all", new_fused_kv_head_fused_kv_fuse_mode: str = "all", kv_fusion: "torch.nn.Module | None" = None, ): """Load Wan components and build a validated decoupled FasterWAM model.""" if video_dit_config is None: raise ValueError( "`video_dit_config` is required for " "FasterWAMDecoupled.from_wan22_pretrained()." ) if "text_dim" not in video_dit_config: raise ValueError("`video_dit_config['text_dim']` is required.") if not decoupled: raise ValueError("FasterWAMDecoupled requires decoupled=True.") new_fused_kv_projection_mode = _validate_new_fused_kv_projection_mode( kv_source_mode, new_fused_kv_projection_mode, ) components = load_wan22_ti2v_5b_components( device=device, torch_dtype=torch_dtype, model_id=model_id, tokenizer_model_id=tokenizer_model_id, tokenizer_max_len=tokenizer_max_len, redirect_common_files=redirect_common_files, dit_config=video_dit_config, skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, load_text_encoder=load_text_encoder, ) video_expert = components.dit layer_init_mapping = kv_source_mapping # Both "fused_kv" and "fused_mlp" fuse across all/many video layers, so # neither has a 1:1 action->video mapping usable for selective init. When # action_dit_pretrained_path is None (the fused_mlp training case), this # mapping is ignored by ActionDiT.from_pretrained (random init) anyway. if kv_source_mode in ("fused_kv", "new_fused_kv", "fused_mlp"): from .mot_decoupled import compute_kv_source_mapping action_config_for_init = action_dit_config or {} action_num_layers = int(action_config_for_init.get("num_layers", 5)) if kv_source_mapping is not None and len(kv_source_mapping) == action_num_layers: layer_init_mapping = kv_source_mapping else: layer_init_mapping = compute_kv_source_mapping( mode="uniform_end", video_num_layers=len(video_expert.blocks), action_num_layers=action_num_layers, ) action_expert = ActionDiT.from_pretrained( action_dit_config=action_dit_config, action_dit_pretrained_path=action_dit_pretrained_path, skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, device=device, torch_dtype=torch_dtype, layer_init_mapping=layer_init_mapping, ) if int(action_expert.num_heads) != int(video_expert.num_heads): raise ValueError( "ActionDiT `num_heads` must match video expert for mixed attention." ) if int(action_expert.attn_head_dim) != int(video_expert.attn_head_dim): raise ValueError( "ActionDiT `attn_head_dim` must match video expert for mixed attention." ) mot_cls = ( MoTDecoupledActionAlignedVideoRoPE if fixed_rope or kv_source_mode == "new_fused_kv" else MoTDecoupled ) logger.info( "FasterWAMDecoupled fixed_rope=%s, kv_source_mode=%s, " "new_fused_kv_rope_mode=%s, aligned_3d_anchor_layout=%s, " "new_fused_kv_projection_mode=%s, " "simple_head_softmax_fuse_mode=%s, " "head_fused_kv_fuse_mode=%s, " "using MoT class %s", fixed_rope, kv_source_mode, new_fused_kv_rope_mode, aligned_3d_action_spatial_anchor_layout, new_fused_kv_projection_mode, new_fused_kv_simple_head_softmax_fuse_mode, new_fused_kv_head_fused_kv_fuse_mode, mot_cls.__name__, ) eef_geometry_identity = None if new_fused_kv_rope_mode in MoTDecoupledActionAlignedVideoRoPE.ee_rope_modes: if eef_calibration_path is None: raise ValueError( f"new_fused_kv_rope_mode={new_fused_kv_rope_mode!r} requires " "eef_calibration_path" ) from fastwam.geometry import ( CAMERA_ORDER, EEF_PROJECTION_VERSION, VAE_SPATIAL_FACTOR, calibration_digest, ) if eef_raw_source_resolution is None: raise ValueError( f"{new_fused_kv_rope_mode} requires " "eef_raw_source_resolution derived from data.shape_meta" ) if eef_processed_video_size is None or len(eef_processed_video_size) != 2: raise ValueError( f"{new_fused_kv_rope_mode} requires eef_processed_video_size " "[H,W] derived from data.video_size" ) processed_h, processed_w = map(int, eef_processed_video_size) patch = tuple(int(v) for v in video_dit_config.get("patch_size", ())) if len(patch) != 3: raise ValueError( "video_dit_config.patch_size must be [T,H,W] to derive the " f"EEF token grid, got {patch}" ) divisor_h = VAE_SPATIAL_FACTOR * patch[1] divisor_w = VAE_SPATIAL_FACTOR * patch[2] if processed_h % divisor_h or processed_w % divisor_w: raise ValueError( f"processed video {processed_h}x{processed_w} is not divisible " f"by VAE*patch {divisor_h}x{divisor_w}" ) eef_geometry_identity = { "calibration_digest": calibration_digest(eef_calibration_path), "raw_source_resolution": int(eef_raw_source_resolution), "token_grid_h": processed_h // divisor_h, "token_grid_w": processed_w // divisor_w, "camera_order": list(CAMERA_ORDER), "projection_version": EEF_PROJECTION_VERSION, } mot_kwargs = {} if mot_cls is MoTDecoupledActionAlignedVideoRoPE: mot_kwargs["new_fused_kv_rope_mode"] = new_fused_kv_rope_mode mot_kwargs["eef_geometry_identity"] = eef_geometry_identity mot_kwargs["aligned_3d_action_spatial_anchor_layout"] = ( aligned_3d_action_spatial_anchor_layout ) mot_kwargs["new_fused_kv_projection_mode"] = new_fused_kv_projection_mode if kv_source_mode == "new_fused_kv": mot_kwargs["new_fused_kv_pos_embed_max_tokens"] = new_fused_kv_pos_embed_max_tokens mot_kwargs["new_fused_kv_pos_embed_dim"] = new_fused_kv_pos_embed_dim mot_kwargs["new_fused_kv_mlp_mixer_num_blocks"] = new_fused_kv_mlp_mixer_num_blocks mot_kwargs["new_fused_kv_mlp_mixer_token_mlp_ratio"] = new_fused_kv_mlp_mixer_token_mlp_ratio mot_kwargs["new_fused_kv_mlp_mixer_channel_mlp_ratio"] = new_fused_kv_mlp_mixer_channel_mlp_ratio mot_kwargs["new_fused_kv_head_fused_kv_sin2d_pe_mlp_hidden_dim"] = new_fused_kv_head_fused_kv_sin2d_pe_mlp_hidden_dim mot_kwargs["new_fused_kv_simple_head_softmax_fuse_mode"] = ( new_fused_kv_simple_head_softmax_fuse_mode ) mot_kwargs["new_fused_kv_head_fused_kv_fuse_mode"] = ( new_fused_kv_head_fused_kv_fuse_mode ) mot = mot_cls( mixtures={"video": video_expert, "action": action_expert}, video_num_layers=len(video_expert.blocks), action_num_layers=len(action_expert.blocks), num_heads=int(video_expert.num_heads), attn_head_dim=int(video_expert.attn_head_dim), mot_checkpoint_mixed_attn=mot_checkpoint_mixed_attn, kv_source_mapping=kv_source_mapping, kv_source_mode=kv_source_mode, kv_fusion=kv_fusion, **mot_kwargs, ) model = cls( video_expert=video_expert, action_expert=action_expert, mot=mot, vae=components.vae, text_encoder=components.text_encoder, tokenizer=components.tokenizer, text_dim=int(video_dit_config["text_dim"]), proprio_dim=proprio_dim, device=device, torch_dtype=torch_dtype, video_train_shift=video_train_shift, video_infer_shift=video_infer_shift, video_num_train_timesteps=video_num_train_timesteps, action_train_shift=action_train_shift, action_infer_shift=action_infer_shift, action_num_train_timesteps=action_num_train_timesteps, loss_lambda_video=loss_lambda_video, loss_lambda_action=loss_lambda_action, ) model.eef_calibration_path = eef_calibration_path model.model_paths = { "video_dit": components.dit_path, "vae": components.vae_path, "text_encoder": components.text_encoder_path, "tokenizer": components.tokenizer_path, "action_dit_backbone": ( "SKIPPED_PRETRAIN" if skip_dit_load_from_pretrain else action_dit_pretrained_path ), } return model def create_fasterwam_decoupled( model_id: str, tokenizer_model_id: str, tokenizer_max_len: int, load_text_encoder: bool, proprio_dim: int | None, redirect_common_files: bool, mot_checkpoint_mixed_attn: bool, action_dit_pretrained_path: str | None, skip_dit_load_from_pretrain: bool, decoupled: bool, kv_source_mode: str, video_dit_config: dict, action_dit_config: dict, video_scheduler: dict, action_scheduler: dict, loss: dict | None = None, device: str = "cuda", model_dtype: torch.dtype = torch.bfloat16, fixed_rope: bool = True, new_fused_kv_rope_mode: str = "aligned_3d", aligned_3d_action_spatial_anchor_layout: str | None = "center", eef_calibration_path: str | None = None, eef_raw_source_resolution: int | None = None, eef_processed_video_size: list[int] | tuple[int, int] | None = None, new_fused_kv_projection_mode: str = "full", new_fused_kv_pos_embed_max_tokens: int = 4096, new_fused_kv_pos_embed_dim: int = 128, new_fused_kv_mlp_mixer_num_blocks: int = 1, new_fused_kv_mlp_mixer_token_mlp_ratio: float = 4.0, new_fused_kv_mlp_mixer_channel_mlp_ratio: float = 4.0, new_fused_kv_head_fused_kv_sin2d_pe_mlp_hidden_dim: int = 512, new_fused_kv_simple_head_softmax_fuse_mode: str = "all", new_fused_kv_head_fused_kv_fuse_mode: str = "all", fusion_hidden_dim: int = 64, fusion_use_norm: bool = True, ) -> FasterWAMDecoupled: """Hydra factory for ``FasterWAMDecoupled``.""" from omegaconf import DictConfig, OmegaConf from .mot_decoupled import compute_kv_source_mapping if isinstance(video_dit_config, DictConfig): video_dit_config = OmegaConf.to_container(video_dit_config, resolve=True) if isinstance(action_dit_config, DictConfig): action_dit_config = OmegaConf.to_container(action_dit_config, resolve=True) if isinstance(video_scheduler, DictConfig): video_scheduler = OmegaConf.to_container(video_scheduler, resolve=True) if isinstance(action_scheduler, DictConfig): action_scheduler = OmegaConf.to_container(action_scheduler, resolve=True) if isinstance(loss, DictConfig): loss = OmegaConf.to_container(loss, resolve=True) if not isinstance(video_dit_config, dict): raise ValueError(f"`video_dit_config` must resolve to dict, got {type(video_dit_config)}") if not isinstance(action_dit_config, dict): raise ValueError(f"`action_dit_config` must resolve to dict, got {type(action_dit_config)}") if not isinstance(video_scheduler, dict): raise ValueError(f"`video_scheduler` must resolve to dict, got {type(video_scheduler)}") if not isinstance(action_scheduler, dict): raise ValueError(f"`action_scheduler` must resolve to dict, got {type(action_scheduler)}") if loss is None: loss = {} if not isinstance(loss, dict): raise ValueError(f"`loss` must resolve to dict, got {type(loss)}") if not decoupled: raise ValueError("create_fasterwam_decoupled requires decoupled=True.") new_fused_kv_projection_mode = _validate_new_fused_kv_projection_mode( kv_source_mode, new_fused_kv_projection_mode, ) required_action_scheduler_keys = {"train_shift", "infer_shift", "num_train_timesteps"} missing_keys = required_action_scheduler_keys - set(action_scheduler.keys()) if missing_keys: raise ValueError( f"`action_scheduler` missing required keys: {sorted(missing_keys)}" ) action_dit_config = dict(action_dit_config) kv_fusion = None # set by the fused_mlp branch below layer_selected = action_dit_config.pop("layer_selected", None) if layer_selected is not None: # layer_selected directly selects video KV sources, which is mutually # exclusive with the learned MLP fusion of fused_mlp mode. if kv_source_mode == "fused_mlp": raise ValueError( "layer_selected and kv_source_mode='fused_mlp' are incompatible" ) action_dit_config["num_layers"] = len(layer_selected) kv_source_mapping = list(layer_selected) logger.info( "layer_selected=%s, overriding num_layers=%d, kv_source_mapping=%s", layer_selected, len(layer_selected), kv_source_mapping, ) elif kv_source_mode == "fused_mlp": # Build a runtime KVFusionModule that fuses all N video layers' K/V into # per-action-layer K/V via a learned MLP. Runtime routing is handled by # this module; the mapping below is only used for weight init selection. from .kv_fusion import KVFusionModule num_video_layers = int(video_dit_config.get("num_layers", 30)) num_action_layers = int(action_dit_config.get("num_layers", 5)) num_heads = int(video_dit_config.get("num_heads", 24)) attn_head_dim = int(video_dit_config.get("attn_head_dim", 128)) init_mode = "uniform_end" if num_action_layers <= num_video_layers else "final_only" kv_source_mapping = compute_kv_source_mapping( mode=init_mode, video_num_layers=num_video_layers, action_num_layers=num_action_layers, ) kv_fusion = KVFusionModule( num_action_layers=num_action_layers, num_video_layers=num_video_layers, attn_hidden_dim=num_heads * attn_head_dim, fusion_hidden_dim=fusion_hidden_dim, fusion_use_norm=fusion_use_norm, dtype=model_dtype, ) logger.info( "KV source mode: fused_mlp (all %d video layers fused via MLP), " "fusion params: %.1fK", num_video_layers, sum(p.numel() for p in kv_fusion.parameters()) / 1e3, ) else: kv_source_mapping = compute_kv_source_mapping( mode=kv_source_mode, video_num_layers=int(video_dit_config.get("num_layers", 30)), action_num_layers=int(action_dit_config.get("num_layers", 5)), ) logger.info("KV source mode: %s, mapping: %s", kv_source_mode, kv_source_mapping) return FasterWAMDecoupled.from_wan22_pretrained( device=device, torch_dtype=model_dtype, model_id=model_id, tokenizer_model_id=tokenizer_model_id, tokenizer_max_len=int(tokenizer_max_len), load_text_encoder=bool(load_text_encoder), proprio_dim=(None if proprio_dim is None else int(proprio_dim)), redirect_common_files=bool(redirect_common_files), video_dit_config=video_dit_config, action_dit_config=action_dit_config, action_dit_pretrained_path=action_dit_pretrained_path, skip_dit_load_from_pretrain=bool(skip_dit_load_from_pretrain), mot_checkpoint_mixed_attn=bool(mot_checkpoint_mixed_attn), video_train_shift=float(video_scheduler.get("train_shift", 5.0)), video_infer_shift=float(video_scheduler.get("infer_shift", 5.0)), video_num_train_timesteps=int(video_scheduler.get("num_train_timesteps", 1000)), action_train_shift=float(action_scheduler["train_shift"]), action_infer_shift=float(action_scheduler["infer_shift"]), action_num_train_timesteps=int(action_scheduler["num_train_timesteps"]), loss_lambda_video=float(loss.get("lambda_video", 1.0)), loss_lambda_action=float(loss.get("lambda_action", 1.0)), decoupled=True, kv_source_mapping=kv_source_mapping, kv_source_mode=kv_source_mode, fixed_rope=bool(fixed_rope), new_fused_kv_rope_mode=str(new_fused_kv_rope_mode), aligned_3d_action_spatial_anchor_layout=( None if aligned_3d_action_spatial_anchor_layout is None else str(aligned_3d_action_spatial_anchor_layout) ), eef_calibration_path=eef_calibration_path, eef_raw_source_resolution=( None if eef_raw_source_resolution is None else int(eef_raw_source_resolution) ), eef_processed_video_size=eef_processed_video_size, new_fused_kv_projection_mode=new_fused_kv_projection_mode, new_fused_kv_pos_embed_max_tokens=int(new_fused_kv_pos_embed_max_tokens), new_fused_kv_pos_embed_dim=int(new_fused_kv_pos_embed_dim), new_fused_kv_mlp_mixer_num_blocks=int(new_fused_kv_mlp_mixer_num_blocks), new_fused_kv_mlp_mixer_token_mlp_ratio=float(new_fused_kv_mlp_mixer_token_mlp_ratio), new_fused_kv_mlp_mixer_channel_mlp_ratio=float(new_fused_kv_mlp_mixer_channel_mlp_ratio), new_fused_kv_head_fused_kv_sin2d_pe_mlp_hidden_dim=int( new_fused_kv_head_fused_kv_sin2d_pe_mlp_hidden_dim ), new_fused_kv_simple_head_softmax_fuse_mode=str( new_fused_kv_simple_head_softmax_fuse_mode ), new_fused_kv_head_fused_kv_fuse_mode=str( new_fused_kv_head_fused_kv_fuse_mode ), kv_fusion=kv_fusion, )