Blame for multics/models/blocks.py

import torch
import torch.nn as nn
import math
from typing import Tuple, Optional
from einops import rearrange, repeat
from .rope import (
    modulate,
    precompute_freqs_cis,
    rope_apply,
)
from .hier_masks import (
    _squeeze_hier_batch,
    _slice_hier_active_group_for_cached_mask,
    _refuse_membership_on_scalar_kernel,
    build_stage2_hier_self_attn_mask,
    build_stage2_hier_flex_block_mask,
    get_temporal_block_causal_mask,
    BlockMask,
    FLEX_ATTENTION_AVAILABLE,
)
from .attention import (
    flash_attention,
    _rolling_cache_append,
    SelfAttention,
    CrossAttention,
)

_HIER_MASK_DEBUG_DUMPED = False


class GateModule(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, gate, residual):
        return x + gate * residual


class DiTBlock(nn.Module):
    def __init__(
        self,
        has_image_input: bool,
        dim: int,
        num_heads: int,
        ffn_dim: int,
        eps: float = 1e-06,
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.ffn_dim = ffn_dim
        self.self_attn = SelfAttention(dim, num_heads, eps)
        self.cross_attn = CrossAttention(
            dim, num_heads, eps, has_image_input=has_image_input
        )
        self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
        self.norm3 = nn.LayerNorm(dim, eps=eps)
        self.ffn = nn.Sequential(
            nn.Linear(dim, ffn_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(ffn_dim, dim),
        )
        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
        self.gate = GateModule()
        self.attention_mode = "dense"
        self.attn_impl = "sdpa"
        self.num_hub_tokens = 8
        self.num_groups = 1
        self.num_local_hub_tokens = 0
        self.num_global_hub_tokens = 8
        self.hier_no_local_hub = False
        self.hier_direct_enabled = False
        self.num_frame_per_block = 7
        self.local_attn_size = 24
        self.hub_embed = None

    def _is_sparse_hub(self):
        return True

    def _modulate_hub(self, x_hub, t_mod, has_seq, chunk_dim):
        k_len = self.num_hub_tokens
        batch_size = x_hub.shape[0]
        num_frames = x_hub.shape[1] // k_len
        assert x_hub.shape[1] == num_frames * k_len, (
            f"x_hub 序列长度 {x_hub.shape[1]} 不是 K={k_len} 的整数倍"
        )
        hub = self.norm1(x_hub)
        mod = self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod
        shift_msa, scale_msa, gate_msa, _, _, _ = mod.chunk(6, dim=chunk_dim)
        if has_seq:
            shift_h = shift_msa.expand(-1, -1, k_len, -1).reshape(
                batch_size, num_frames * k_len, self.dim
            )
            scale_h = scale_msa.expand(-1, -1, k_len, -1).reshape(
                batch_size, num_frames * k_len, self.dim
            )
            gate_h = gate_msa.expand(-1, -1, k_len, -1).reshape(
                batch_size, num_frames * k_len, self.dim
            )
        else:
            shift_h = shift_msa.unsqueeze(1).expand(-1, num_frames * k_len, -1)
            scale_h = scale_msa.unsqueeze(1).expand(-1, num_frames * k_len, -1)
            gate_h = gate_msa.unsqueeze(1).expand(-1, num_frames * k_len, -1)
        hub_in = modulate(hub, shift_h, scale_h)
        return (hub_in, gate_h)

    def _build_camera_tokens(
        self,
        camera_pose,
        num_frames: int,
        num_patches: int,
        num_players: int,
        ref_x: torch.Tensor,
        include_hub: bool,
    ):
        """Expand per-player per-frame camera12 to self-attn token order.

        camera_pose accepts [P, F, C] (current B=1 path) or [B, P, F, C].
        Output follows the same player-major order as input_x after
        rearrange("(b p) l d -> b (p l) d", p=P): [B, P*F*L, C],
        optionally followed by zero camera tokens for the hub band [B, F*K, C].
        """
        if camera_pose is None:
            return None
        cam = camera_pose
        if cam.dim() == 3:
            cam = cam.unsqueeze(0)
        elif cam.dim() != 4:
            raise ValueError(
                f"camera_pose must be [P,F,C] or [B,P,F,C], got shape {tuple(camera_pose.shape)}"
            )
        if cam.shape[1] != num_players or cam.shape[2] != num_frames:
            raise ValueError(
                f"camera_pose expects P={num_players}, F={num_frames}, got {tuple(cam.shape[1:3])}"
            )
        if cam.shape[0] != ref_x.shape[0]:
            raise ValueError(
                f"camera_pose batch {cam.shape[0]} must match self-attn batch {ref_x.shape[0]}"
            )
        cam = cam.to(device=ref_x.device, dtype=ref_x.dtype)
        b, p, f, c = cam.shape
        cam_agent = cam[:, :, :, None, :].expand(b, p, f, num_patches, c)
        cam_agent = cam_agent.reshape(b, p * f * num_patches, c)
        if include_hub:
            cam_hub = cam_agent.new_zeros(b, f * self.num_hub_tokens, c)
            cam_agent = torch.cat([cam_agent, cam_hub], dim=1)
        return cam_agent

    def _self_attn_with_hub(
        self,
        input_x,
        hub_in,
        freqs,
        mask=None,
        kv_cache=None,
        grid_sizes=None,
        seq_lens=None,
        current_start=0,
        cache_start=None,
        num_frames=None,
        num_players=None,
        cam_tokens=None,
        group_labels=None,
        active_group=None,
        direct_player=False,
        prebuilt_hier_tiled_layouts=None,
        membership=None,
        global_only=None,
    ):
        seq = torch.cat([input_x, hub_in], dim=1)
        agent_len = input_x.shape[1]
        if kv_cache is None:
            if getattr(self, "attn_impl", "sdpa") == "hier_tiled_ste":
                if membership is not None:
                    _refuse_membership_on_scalar_kernel(attn_impl="hier_tiled_ste")
                if group_labels is None or active_group is None:
                    raise ValueError(
                        "hier_tiled_ste requires explicit group_labels and active_group"
                    )
                if num_frames is None or num_players is None:
                    raise ValueError(
                        "hier_tiled_ste requires num_frames and num_players"
                    )
                num_patches = input_x.shape[1] // (int(num_players) * int(num_frames))
                out = self.self_attn.forward_hier_tiled_ste(
                    seq,
                    freqs,
                    mask,
                    group_labels=group_labels,
                    active_group=active_group,
                    num_players=int(num_players),
                    num_frames=int(num_frames),
                    num_patches=int(num_patches),
                    num_groups=self.num_groups,
                    num_local_hub_tokens=self.num_local_hub_tokens,
                    num_global_hub_tokens=self.num_global_hub_tokens,
                    direct_player=bool(direct_player),
                    tiled_layouts=prebuilt_hier_tiled_layouts,
                    cam_tokens=cam_tokens,
                )
            else:
                out = self.self_attn(seq, freqs, mask, cam_tokens=cam_tokens)
        else:
            out = self.self_attn.causal_forward(
                seq,
                freqs,
                mask=None,
                grid_sizes=grid_sizes,
                seq_lens=seq_lens,
                kv_cache=kv_cache,
                current_start=current_start,
                cache_start=cache_start,
                num_frames=num_frames,
                num_players=num_players,
                num_hub_tokens=self.num_hub_tokens,
                cam_tokens=cam_tokens,
                group_labels=group_labels,
                active_group=active_group,
                direct_player=direct_player,
                membership=membership,
                global_only=global_only,
            )
        assert out.shape[1] == agent_len + hub_in.shape[1], (
            f"self-attn 全序列长度 {out.shape[1]} != agent {agent_len} + hub {hub_in.shape[1]}"
        )
        return (out[:, :agent_len], out[:, agent_len:])

    @staticmethod
    def _ensure_action_cache(
        parent, name, batch, n_heads, head_dim, tokens_per_frame, frames, device, dtype
    ):
        """Lazily allocate (or re-allocate on shape change) a rolling action KV cache
        stored inside the per-layer kv_cache dict under `name`."""
        size = frames * tokens_per_frame
        c = parent.get(name)
        if (
            c is None
            or c["k"].shape[0] != batch
            or c["k"].shape[1] < size
            or (c["k"].shape[2] != n_heads)
            or (c["k"].shape[3] != head_dim)
        ):
            parent[name] = {
                "k": torch.zeros(
                    [batch, size, n_heads, head_dim], device=device, dtype=dtype
                ),
                "v": torch.zeros(
                    [batch, size, n_heads, head_dim], device=device, dtype=dtype
                ),
                "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
                "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
            }
        return parent[name]

    def _action_temporal_cached_attn(
        self, attn, x, freqs_temporal, kv_cache, start_frame, n_block, local_attn
    ):
        """Cross-block KV-cached temporal self-attention for the mouse/action branch.

        Mirrors the training-time block-causal+window mask (`get_temporal_block_causal_mask`)
        but feeds historical frames from a rolling cache instead of the current block only.
        x: [B, f, D] with the sequence axis = frames (1 token per frame)."""
        B, f, D = x.shape
        nh = attn.num_heads
        hd = D // nh
        q = attn.norm_q(attn.q(x))
        k = attn.norm_k(attn.k(x))
        v = attn.v(x)
        roped_q = rope_apply(q, freqs_temporal, nh)
        roped_k = rope_apply(k, freqs_temporal, nh)
        rk = roped_k.view(B, f, nh, hd)
        vv = v.view(B, f, nh, hd)
        retention = local_attn + f if local_attn != -1 else None
        cache = self._ensure_action_cache(
            kv_cache,
            "action_temporal",
            B,
            nh,
            hd,
            tokens_per_frame=1,
            frames=retention if retention is not None else start_frame + f,
            device=x.device,
            dtype=x.dtype,
        )
        attn_start, local_end = _rolling_cache_append(
            cache,
            rk,
            vv,
            current_start=start_frame,
            current_end=start_frame + f,
            local_attn_size=retention if retention is not None else -1,
            sink_size=0,
            tokens_per_frame=1,
        )
        cached_global_start = (
            cache["global_end_index"].item() - cache["local_end_index"].item()
        )
        win = local_end - attn_start
        win_k = cache["k"][:, attn_start:local_end].reshape(B, win, D)
        win_v = cache["v"][:, attn_start:local_end].reshape(B, win, D)
        gq = start_frame + torch.arange(f, device=x.device)
        gk = cached_global_start + torch.arange(attn_start, local_end, device=x.device)
        mask = gk[None, :] // n_block <= gq[:, None] // n_block
        if local_attn != -1:
            mask = mask & (gq[:, None] - gk[None, :] < local_attn)
        out = flash_attention(roped_q, win_k, win_v, num_heads=nh, mask=mask)
        return attn.o(out)

    def _action_keyboard_cached_attn(
        self,
        attn,
        x,
        kbd_ctx,
        kv_cache,
        start_frame,
        n_block,
        local_attn,
        num_patches,
        tokens_per_kbd_frame,
    ):
        """Cross-block KV-cached keyboard cross-attention.

        Video query tokens (current block) attend to cached + current keyboard tokens
        under the same block-causal+window relation used in training."""
        B = x.shape[0]
        D = x.shape[-1]
        nh = attn.num_heads
        hd = D // nh
        f = kbd_ctx.shape[1] // tokens_per_kbd_frame
        ctx = kbd_ctx[:, 257:] if getattr(attn, "has_image_input", False) else kbd_ctx
        q = attn.norm_q(attn.q(x))
        k = attn.norm_k(attn.k(ctx))
        v = attn.v(ctx)
        k4 = k.view(B, f * tokens_per_kbd_frame, nh, hd)
        v4 = v.view(B, f * tokens_per_kbd_frame, nh, hd)
        retention = local_attn + f if local_attn != -1 else None
        cache = self._ensure_action_cache(
            kv_cache,
            "action_keyboard",
            B,
            nh,
            hd,
            tokens_per_frame=tokens_per_kbd_frame,
            frames=retention if retention is not None else start_frame + f,
            device=x.device,
            dtype=x.dtype,
        )
        attn_start, local_end = _rolling_cache_append(
            cache,
            k4,
            v4,
            current_start=start_frame * tokens_per_kbd_frame,
            current_end=(start_frame + f) * tokens_per_kbd_frame,
            local_attn_size=retention if retention is not None else -1,
            sink_size=0,
            tokens_per_frame=tokens_per_kbd_frame,
        )
        cached_global_start = (
            cache["global_end_index"].item() - cache["local_end_index"].item()
        )
        win = local_end - attn_start
        win_k = cache["k"][:, attn_start:local_end].reshape(B, win, D)
        win_v = cache["v"][:, attn_start:local_end].reshape(B, win, D)
        gq = start_frame + torch.arange(f, device=x.device)
        key_tok = cached_global_start + torch.arange(
            attn_start, local_end, device=x.device
        )
        gk = key_tok // tokens_per_kbd_frame
        frame_mask = gk[None, :] // n_block <= gq[:, None] // n_block
        if local_attn != -1:
            frame_mask = frame_mask & (gq[:, None] - gk[None, :] < local_attn)
        mask = frame_mask.repeat_interleave(num_patches, dim=0)
        out = flash_attention(q, win_k, win_v, num_heads=nh, mask=mask)
        return attn.o(out)

    def _get_action_masks(
        self,
        num_frames,
        num_patches,
        num_players,
        loss_type,
        device,
        prebuilt_self_attn_block_mask=None,
    ):
        local_attn = self.local_attn_size
        n_block = self.num_frame_per_block
        group_labels = getattr(self, "_hier_group_labels", None)
        active_group = getattr(self, "_hier_active_group", None)
        membership = getattr(self, "_hier_membership", None)
        global_only = getattr(self, "_hier_global_only", None)
        if group_labels is None and membership is None:
            raise ValueError(
                "hier_hub requires block._hier_group_labels [T,P] or _hier_membership [T,N,P]"
            )
        direct_enabled = bool(getattr(self, "hier_direct_enabled", False))
        if prebuilt_self_attn_block_mask is not None:
            if not isinstance(prebuilt_self_attn_block_mask, BlockMask):
                raise TypeError(
                    "prebuilt_self_attn_block_mask must be a FlexAttention BlockMask"
                )
            block_mask = prebuilt_self_attn_block_mask
        elif (
            getattr(self, "attn_impl", "sdpa") in ("flex", "hier_tiled_ste")
            and FLEX_ATTENTION_AVAILABLE
        ):
            block_mask, _, _ = build_stage2_hier_flex_block_mask(
                num_players,
                num_frames,
                num_patches,
                self.num_groups,
                self.num_local_hub_tokens,
                self.num_global_hub_tokens,
                n_block,
                local_attn,
                group_labels,
                active_group,
                direct_enabled,
                device,
                membership=membership,
                global_only=global_only,
                hier_no_local_hub=bool(getattr(self, "hier_no_local_hub", False)),
            )
        else:
            block_mask = build_stage2_hier_self_attn_mask(
                num_players,
                num_frames,
                num_patches,
                self.num_groups,
                self.num_local_hub_tokens,
                self.num_global_hub_tokens,
                n_block,
                local_attn,
                group_labels,
                active_group,
                direct_enabled,
                device=device,
                membership=membership,
                global_only=global_only,
                hier_no_local_hub=bool(getattr(self, "hier_no_local_hub", False)),
            )
        frame_mask = get_temporal_block_causal_mask(
            num_frames, n_block, local_attn, device=device
        )
        block_mask_mouse = frame_mask
        block_mask_keyboard = repeat(
            frame_mask, "f g -> (f q) (g k)", q=num_patches, k=12
        )
        import os

        DEBUG = os.environ.get("STAGE2_MASK_DEBUG", "0") not in (
            "0",
            "",
            "false",
            "False",
        )
        if DEBUG and os.environ.get("STAGE2_MASK_DEBUG_ONCE", "1") not in (
            "0",
            "",
            "false",
            "False",
        ):
            global _HIER_MASK_DEBUG_DUMPED
            if _HIER_MASK_DEBUG_DUMPED:
                DEBUG = False
            else:
                _HIER_MASK_DEBUG_DUMPED = True
        if DEBUG:
            import numpy as np
            import imageio

            _dir = os.environ.get("STAGE2_MASK_DEBUG_DIR", ".")
            os.makedirs(_dir, exist_ok=True)
            if torch.is_tensor(block_mask):
                _self_attn_dense = block_mask
            else:
                _self_attn_dense = build_stage2_hier_self_attn_mask(
                    num_players,
                    num_frames,
                    num_patches,
                    self.num_groups,
                    self.num_local_hub_tokens,
                    self.num_global_hub_tokens,
                    n_block,
                    local_attn,
                    group_labels,
                    active_group,
                    direct_enabled,
                    device=device,
                    membership=membership,
                    global_only=global_only,
                    hier_no_local_hub=bool(getattr(self, "hier_no_local_hub", False)),
                )
            _self_attn_cpu = _self_attn_dense.detach().cpu().bool()

            def _write_mask_image(_name, _m, _size=1024):
                if torch.is_tensor(_m):
                    _arr = _m.detach().cpu().numpy()
                else:
                    _arr = np.asarray(_m)
                if _arr.dtype == np.bool_:
                    _arr = _arr.astype(np.uint8) * 255
                elif _arr.dtype != np.uint8:
                    _arr = np.uint8(255.0 * _arr.astype(np.float32))
                try:
                    import cv2

                    _arr = cv2.resize(
                        _arr, (_size, _size), interpolation=cv2.INTER_NEAREST
                    )
                except Exception:
                    pass
                imageio.imwrite(os.path.join(_dir, "mask_%s.jpg" % _name), _arr)

            _dbg = [
                ("self_attn", _self_attn_cpu),
                ("mouse", block_mask_mouse),
                ("keyboard", block_mask_keyboard),
            ]
            for _name, _m in _dbg:
                _write_mask_image(_name, _m)
            _agent_len = int(num_players) * int(num_frames) * int(num_patches)
            _seq_len = int(_self_attn_cpu.shape[0])
            _hub_len = max(0, _seq_len - _agent_len)
            if _hub_len > 0:
                _crop_env = os.environ.get("STAGE2_MASK_DEBUG_CROP_AGENT_TOKENS")
                if _crop_env:
                    _crop_agent = int(_crop_env)
                else:
                    _crop_agent = min(
                        _agent_len,
                        max(
                            4 * _hub_len,
                            int(num_players) * int(num_patches) * max(1, int(n_block)),
                            512,
                        ),
                        4096,
                    )
                _crop_agent = max(0, min(_agent_len, _crop_agent))
                _crop_start = max(0, _agent_len - _crop_agent)
                _write_mask_image(
                    "self_attn_agent_hub_crop",
                    _self_attn_cpu[_crop_start:_seq_len, _crop_start:_seq_len],
                )
            _units = []
            _bounds = [0]
            _T = int(num_frames)
            _L = int(num_patches)
            _P = int(num_players)
            for _p in range(_P):
                _p_base = _p * _T * _L
                for _t in range(_T):
                    _units.append((_p_base + _t * _L, _p_base + (_t + 1) * _L))
                _bounds.append(len(_units))
            _N = int(self.num_groups)
            _KL = int(self.num_local_hub_tokens)
            _KG = int(self.num_global_hub_tokens)
            _local_start = _agent_len
            for _g in range(_N):
                for _t in range(_T):
                    _s = _local_start + (_t * _N + _g) * _KL
                    _units.append((_s, _s + _KL))
                _bounds.append(len(_units))
            _global_start = _local_start + _T * _N * _KL
            if _KG > 0:
                for _t in range(_T):
                    _s = _global_start + _t * _KG
                    _units.append((_s, _s + _KG))
                _bounds.append(len(_units))
            _M = len(_units)
            _coarse = np.zeros((_M, _M), dtype=np.uint8)
            for _qi, (_qs, _qe) in enumerate(_units):
                if _qe <= _qs:
                    continue
                for _ki, (_ks, _ke) in enumerate(_units):
                    if _ke <= _ks:
                        continue
                    if bool(_self_attn_cpu[_qs:_qe, _ks:_ke].any().item()):
                        _coarse[_qi, _ki] = 255
            for _b in _bounds:
                if 0 <= _b < _M:
                    _coarse[_b, :] = np.maximum(_coarse[_b, :], 127)
                    _coarse[:, _b] = np.maximum(_coarse[:, _b], 127)
            _write_mask_image("self_attn_coarse", _coarse)
        return (block_mask, block_mask_mouse, block_mask_keyboard)

    def forward(
        self,
        x,
        context,
        mouse_BPFD,
        keyboard_BPFD,
        t_mod,
        freqs,
        loss_type="bidirectional",
        x_hub=None,
        grid_hw=None,
        camera_pose: Optional[torch.Tensor] = None,
        prebuilt_self_attn_block_mask=None,
        prebuilt_hier_tiled_layouts=None,
    ):
        has_seq = len(t_mod.shape) == 4
        chunk_dim = 2 if has_seq else 1
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod
        ).chunk(6, dim=chunk_dim)
        if has_seq:
            hw = x.shape[1] // t_mod.shape[1]
            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [
                v.expand(-1, -1, hw, -1).reshape(v.shape[0], -1, v.shape[-1])
                for v in [
                    shift_msa,
                    scale_msa,
                    gate_msa,
                    shift_mlp,
                    scale_mlp,
                    gate_mlp,
                ]
            ]
        num_frames = keyboard_BPFD.shape[1]
        num_players = keyboard_BPFD.shape[0]
        num_patches = x.shape[1] // num_frames
        block_mask, block_mask_mouse, block_mask_keyboard = self._get_action_masks(
            num_frames,
            num_patches,
            num_players,
            "diffusion_forcing",
            x.device,
            prebuilt_self_attn_block_mask=prebuilt_self_attn_block_mask,
        )
        x_in = x
        delta_q_bp, delta_o_bp, dq_agent = (None, None, None)
        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
        input_x = rearrange(input_x, "(b p) l d -> b (p l) d", p=num_players)
        cam_tokens = self._build_camera_tokens(
            camera_pose,
            num_frames,
            num_patches,
            num_players,
            input_x,
            include_hub=self._is_sparse_hub(),
        )
        assert x_hub is not None, "hier_hub requires an external x_hub residual stream"
        t_hub = t_mod[:1] if has_seq and t_mod.shape[0] == num_players else t_mod
        hub_in, hub_gate_msa = self._modulate_hub(x_hub, t_hub, has_seq, chunk_dim)
        self_attn_out, hub_out = self._self_attn_with_hub(
            input_x,
            hub_in,
            freqs,
            mask=block_mask,
            num_frames=num_frames,
            num_players=num_players,
            cam_tokens=cam_tokens,
            group_labels=getattr(self, "_hier_group_labels", None),
            active_group=getattr(self, "_hier_active_group", None),
            direct_player=bool(getattr(self, "hier_direct_enabled", False)),
            prebuilt_hier_tiled_layouts=prebuilt_hier_tiled_layouts,
            membership=getattr(self, "_hier_membership", None),
            global_only=getattr(self, "_hier_global_only", None),
        )
        x_hub = self.gate(x_hub, hub_gate_msa, hub_out)
        assert self_attn_out.shape[1] == num_players * num_frames * num_patches, (
            f"self_attn_out agent len {self_attn_out.shape[1]} != P*F*L={num_players}*{num_frames}*{num_patches}"
        )
        self_attn_out = rearrange(
            self_attn_out, "b (p l) d -> (b p) l d", p=num_players
        )
        x = self.gate(x, gate_msa, self_attn_out)
        if delta_o_bp is not None:
            x = x + delta_o_bp
        x = x + self.cross_attn(self.norm3(x), context)
        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
        input_x = torch.cat([input_x, mouse_BPFD], dim=-1)
        input_x = self.action_mlp(input_x)
        bp_dim = input_x.shape[0]
        spatial_tokens_per_frame = input_x.shape[1] // num_frames
        input_x = rearrange(input_x, "b (f l) d -> (b l) f d", f=num_frames)
        n_heads = self.self_attn.num_heads
        head_dim = input_x.shape[-1] // n_heads
        freqs_temporal = precompute_freqs_cis(head_dim, end=input_x.shape[1]).to(
            input_x.device
        )
        freqs_temporal = freqs_temporal.unsqueeze(1)
        action_temporal_attn_output = self.action_temporal_attn(
            input_x, freqs_temporal, block_mask_mouse
        )
        action_temporal_attn_output = rearrange(
            action_temporal_attn_output,
            "(b l) f d -> b (f l) d",
            b=bp_dim,
            l=spatial_tokens_per_frame,
        )
        x = x + action_temporal_attn_output
        keyboard_BPFD = rearrange(keyboard_BPFD, "p f l d -> p (f l) d")
        x = x + self.action_cross_attn(
            self.norm3(x), keyboard_BPFD, block_mask_keyboard
        )
        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
        x = self.gate(x, gate_mlp, self.ffn(input_x))
        return (x, x_hub)

    @staticmethod
    def get_causal_attn_mask(
        num_q_blocks, num_k_blocks, q_block_size, k_block_size, sliding_block_size
    ):
        i = torch.arange(num_q_blocks)[:, None]
        j = torch.arange(num_k_blocks)[None, :]
        if sliding_block_size == -1:
            block_mask = j <= i
        else:
            block_mask = (j <= i) & (j >= i - sliding_block_size + 1)
        block_mask = repeat(
            block_mask,
            "q k -> (q q_block) (k k_block)",
            q_block=q_block_size,
            k_block=k_block_size,
        )
        return block_mask

    def causal_forward(
        self,
        x,
        context,
        mouse_BPFD,
        keyboard_BPFD,
        t_mod,
        freqs,
        loss_type="bidirectional",
        grid_sizes=None,
        seq_lens=None,
        kv_cache=None,
        crossattn_cache=None,
        current_start=0,
        cache_start=None,
        x_hub=None,
        camera_pose: Optional[torch.Tensor] = None,
        group_labels: Optional[torch.Tensor] = None,
        active_group: Optional[torch.Tensor] = None,
        membership: Optional[torch.Tensor] = None,
        global_only: Optional[torch.Tensor] = None,
    ):
        has_seq = len(t_mod.shape) == 4
        chunk_dim = 2 if has_seq else 1
        x_in = x
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod
        ).chunk(6, dim=chunk_dim)
        if has_seq:
            hw = x.shape[1] // t_mod.shape[1]
            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [
                v.expand(-1, -1, hw, -1).reshape(v.shape[0], -1, v.shape[-1])
                for v in [
                    shift_msa,
                    scale_msa,
                    gate_msa,
                    shift_mlp,
                    scale_mlp,
                    gate_mlp,
                ]
            ]
        num_frames = keyboard_BPFD.shape[1]
        num_players = keyboard_BPFD.shape[0]
        num_patches = x.shape[1] // num_frames
        start_frame = int(current_start) // max(int(num_players) * int(num_patches), 1)
        self._hier_group_labels = group_labels
        self._hier_active_group = _slice_hier_active_group_for_cached_mask(
            active_group, start_frame, num_frames
        )
        if membership is not None:
            mem = _squeeze_hier_batch(membership, 3, "membership")
            go = (
                _squeeze_hier_batch(global_only, 2, "global_only")
                if global_only is not None
                else None
            )
            if mem.shape[0] > num_frames:
                self._hier_membership = mem[start_frame : start_frame + num_frames]
                self._hier_global_only = (
                    None if go is None else go[start_frame : start_frame + num_frames]
                )
                if active_group is None:
                    self._hier_active_group = mem[
                        start_frame : start_frame + num_frames
                    ].any(dim=-1)
            else:
                self._hier_membership = mem
                self._hier_global_only = go
        else:
            self._hier_membership = None
            self._hier_global_only = None
        block_mask, block_mask_mouse, block_mask_keyboard = self._get_action_masks(
            num_frames, num_patches, num_players, "diffusion_forcing", x.device
        )
        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
        input_x = rearrange(input_x, "(b p) l d -> b (p l) d", p=num_players)
        cam_tokens = self._build_camera_tokens(
            camera_pose,
            num_frames,
            num_patches,
            num_players,
            input_x,
            include_hub=self._is_sparse_hub(),
        )
        if grid_sizes is not None:
            bsz = grid_sizes.shape[0] // num_players
            grid_sizes_bp = rearrange(grid_sizes, "(b p) c -> b p c", p=num_players)
            grid_sizes_repeated = torch.cat(
                [
                    grid_sizes_bp[:, 0, :],
                    torch.full(
                        (bsz, 1),
                        num_players,
                        dtype=grid_sizes.dtype,
                        device=grid_sizes.device,
                    ),
                ],
                dim=1,
            )
        else:
            grid_sizes_repeated = None
        if seq_lens is not None:
            seq_lens_repeated = (
                [sl * num_players for sl in seq_lens] if seq_lens is not None else None
            )
        else:
            seq_lens_repeated = None
        delta_o_bp, dq_agent = (None, None)
        assert x_hub is not None, "hier_hub 需要外部传入 x_hub 残差流"
        t_hub = t_mod[:1] if has_seq and t_mod.shape[0] == num_players else t_mod
        hub_in, hub_gate_msa = self._modulate_hub(x_hub, t_hub, has_seq, chunk_dim)
        self_attn_out, hub_out = self._self_attn_with_hub(
            input_x,
            hub_in,
            freqs,
            mask=block_mask if kv_cache is None else None,
            kv_cache=kv_cache,
            grid_sizes=grid_sizes_repeated,
            seq_lens=seq_lens_repeated,
            current_start=current_start,
            cache_start=cache_start,
            num_frames=num_frames,
            num_players=num_players,
            cam_tokens=cam_tokens,
            group_labels=group_labels,
            active_group=active_group,
            direct_player=bool(getattr(self, "hier_direct_enabled", False)),
            membership=membership,
            global_only=global_only,
        )
        x_hub = self.gate(x_hub, hub_gate_msa, hub_out)
        assert self_attn_out.shape[1] == num_players * num_frames * num_patches, (
            f"self_attn_out agent len {self_attn_out.shape[1]} != P*F*L={num_players}*{num_frames}*{num_patches}"
        )
        self_attn_out = rearrange(
            self_attn_out, "b (p l) d -> (b p) l d", p=num_players
        )
        x = self.gate(x, gate_msa, self_attn_out)
        if delta_o_bp is not None:
            x = x + delta_o_bp
        x = x + self.cross_attn(self.norm3(x), context)
        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
        input_x = torch.cat([input_x, mouse_BPFD], dim=-1)
        input_x = self.action_mlp(input_x)
        bp_dim = input_x.shape[0]
        spatial_tokens_per_frame = input_x.shape[1] // num_frames
        input_x = rearrange(input_x, "b (f l) d -> (b l) f d", f=num_frames)
        n_heads = self.self_attn.num_heads
        head_dim = input_x.shape[-1] // n_heads
        _tpf_allp = num_patches * num_players
        _start_frame_offset = current_start // _tpf_allp if _tpf_allp > 0 else 0
        freqs_temporal = precompute_freqs_cis(
            head_dim, end=_start_frame_offset + input_x.shape[1]
        ).to(input_x.device)
        freqs_temporal = freqs_temporal[
            _start_frame_offset : _start_frame_offset + input_x.shape[1]
        ]
        freqs_temporal = freqs_temporal.unsqueeze(1)
        if kv_cache is not None:
            action_temporal_attn_output = self._action_temporal_cached_attn(
                self.action_temporal_attn,
                input_x,
                freqs_temporal,
                kv_cache,
                _start_frame_offset,
                self.num_frame_per_block,
                self.local_attn_size,
            )
        else:
            action_temporal_attn_output = self.action_temporal_attn(
                input_x, freqs_temporal, block_mask_mouse
            )
        action_temporal_attn_output = rearrange(
            action_temporal_attn_output,
            "(b l) f d -> b (f l) d",
            b=bp_dim,
            l=spatial_tokens_per_frame,
        )
        x = x + action_temporal_attn_output
        keyboard_BPFD = rearrange(keyboard_BPFD, "p f l d -> p (f l) d")
        if kv_cache is not None:
            tokens_per_kbd_frame = keyboard_BPFD.shape[1] // num_frames
            x = x + self._action_keyboard_cached_attn(
                self.action_cross_attn,
                self.norm3(x),
                keyboard_BPFD,
                kv_cache,
                _start_frame_offset,
                self.num_frame_per_block,
                self.local_attn_size,
                num_patches,
                tokens_per_kbd_frame,
            )
        else:
            x = x + self.action_cross_attn(
                self.norm3(x), keyboard_BPFD, block_mask_keyboard
            )
        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
        x = self.gate(x, gate_mlp, self.ffn(input_x))
        return (x, x_hub)


class MLP(torch.nn.Module):
    def __init__(self, in_dim, out_dim, has_pos_emb=False):
        super().__init__()
        self.proj = torch.nn.Sequential(
            nn.LayerNorm(in_dim),
            nn.Linear(in_dim, in_dim),
            nn.GELU(),
            nn.Linear(in_dim, out_dim),
            nn.LayerNorm(out_dim),
        )
        self.has_pos_emb = has_pos_emb
        if has_pos_emb:
            self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280)))

    def forward(self, x):
        if self.has_pos_emb:
            x = x + self.emb_pos.to(dtype=x.dtype, device=x.device)
        return self.proj(x)


class Head(nn.Module):
    def __init__(
        self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float
    ):
        super().__init__()
        self.dim = dim
        self.patch_size = patch_size
        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
        self.head = nn.Linear(dim, out_dim * math.prod(patch_size))
        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)

    def forward(self, x, t_mod):
        if len(t_mod.shape) == 3:
            shift, scale = (
                self.modulation.to(dtype=t_mod.dtype, device=t_mod.device).unsqueeze(1)
                + t_mod.unsqueeze(2)
            ).chunk(2, dim=2)
            shift = shift.squeeze(2)
            scale = scale.squeeze(2)
            hw = x.shape[1] // t_mod.shape[1]
            shift = (
                shift.unsqueeze(2)
                .expand(-1, -1, hw, -1)
                .reshape(shift.shape[0], -1, shift.shape[-1])
            )
            scale = (
                scale.unsqueeze(2)
                .expand(-1, -1, hw, -1)
                .reshape(scale.shape[0], -1, scale.shape[-1])
            )
            x = self.head(self.norm(x) * (1 + scale) + shift)
        else:
            shift, scale = (
                self.modulation.to(dtype=t_mod.dtype, device=t_mod.device)
                + t_mod.unsqueeze(1)
            ).chunk(2, dim=1)
            x = self.head(self.norm(x) * (1 + scale) + shift)
        return x