--- title: model ---

fastvideo.models.dits.matrixgame3.model.DistributedAttention

DistributedAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '', **extra_impl_args)

Distributed attention layer.

fastvideo.models.dits.matrixgame3.model.DistributedAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, original_seq_len: int | None = None, replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, attn_metadata: AttentionMetadata | None = None) -> tuple[torch.Tensor, torch.Tensor | None]

Forward pass for distributed attention.

Parameters

q: torch.Tensor

Query tensor [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor [batch_size, seq_len, num_heads, head_dim]

original_seq_len: int

Original (unpadded) full sequence length

replicated_q: Optional[torch.Tensor]

Replicated query tensor, typically for text tokens

replicated_k: Optional[torch.Tensor]

Replicated key tensor

replicated_v: Optional[torch.Tensor]

Replicated value tensor

Returns

tuple[torch.Tensor, torch.Tensor | None]

Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing: - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided

fastvideo.models.dits.matrixgame3.model.MatrixGame3WanVideoConfig

MatrixGame3WanVideoConfig(arch_config: MatrixGame3WanVideoArchConfig = MatrixGame3WanVideoArchConfig(), prefix: str = 'Wan', quant_config: QuantizationConfig | None = None, _compile_conditions: list = (lambda: [_is_transformer_block])(), *, _resolved_attention_backend: AttentionBackendEnum | None = None)

fastvideo.models.dits.matrixgame3.model.FP32LayerNorm

FP32LayerNorm()

fastvideo.models.dits.matrixgame3.model.FP32LayerNorm.forward

forward(inputs: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.matrixgame3.model.LayerNormScaleShift

LayerNormScaleShift(hidden_size: int, norm_type: str = 'rms', eps: float = 1e-06, elementwise_affine: bool = False, dtype: torch.dtype = torch.float32, compute_dtype: torch.dtype | None = None, prefix: str = '')

Fused operation that combines LayerNorm with scale and shift operations. This reduces memory bandwidth by combining memory-bound operations.

fastvideo.models.dits.matrixgame3.model.LayerNormScaleShift.forward

forward(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor, convert_modulation_dtype: bool = False) -> torch.Tensor

Apply ln followed by scale and shift in a single fused operation.

fastvideo.models.dits.matrixgame3.model.RMSNorm

RMSNorm(hidden_size: int, eps: float = 1e-06, dtype: torch.dtype = torch.float32, var_hidden_size: int | None = None, has_weight: bool = True)

Root mean square normalization.

Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight. Refer to https://arxiv.org/abs/1910.07467

fastvideo.models.dits.matrixgame3.model.RMSNorm.forward_native

forward_native(x: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

PyTorch-native implementation equivalent to forward().

fastvideo.models.dits.matrixgame3.model.RMSNorm.extra_repr

extra_repr() -> str

fastvideo.models.dits.matrixgame3.model.ScaleResidual

ScaleResidual(prefix: str = '')

Applies gated residual connection.

fastvideo.models.dits.matrixgame3.model.ScaleResidual.forward

forward(residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor

Apply gated residual connection.

fastvideo.models.dits.matrixgame3.model.ScaleResidualLayerNormScaleShift

ScaleResidualLayerNormScaleShift(hidden_size: int, norm_type: str = 'rms', eps: float = 1e-06, elementwise_affine: bool = False, dtype: torch.dtype = torch.float32, compute_dtype: torch.dtype | None = None, prefix: str = '', fuse_inference: bool = False)

Fused operation that combines: 1. Gated residual connection 2. LayerNorm 3. Scale and shift operations

This reduces memory bandwidth by combining memory-bound operations.

fastvideo.models.dits.matrixgame3.model.ScaleResidualLayerNormScaleShift.forward

forward(residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor | int, shift: torch.Tensor | None, scale: torch.Tensor | None, convert_modulation_dtype: bool = False) -> tuple[torch.Tensor, torch.Tensor]

Apply gated residual connection, followed by layernorm and scale/shift in a single fused operation.

Pass shift=None, scale=None (both, not just one) to skip the modulation entirely and return the normalized tensor as-is. Callers that need norm-only behavior (e.g. Wan's post-self-attn site) would otherwise pay two full-tensor elementwise passes computing normalized * (1.0 + 0) + 0.

Returns

torch.Tensor

Tuple containing:

torch.Tensor
  • normalized and modulated output
tuple[torch.Tensor, torch.Tensor]
  • residual value (value after residual connection but before normalization)

fastvideo.models.dits.matrixgame3.model.ReplicatedLinear

ReplicatedLinear(input_size: int, output_size: int, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Replicated linear layer.

Parameters

input_size: int

input dimension of the linear layer.

output_size: int

output dimension of the linear layer.

bias: bool

If true, add bias.

skip_bias_add: bool

If true, skip adding bias but instead return it.

params_dtype: torch.dtype | None

Data type for the parameters.

quant_config: QuantizationConfig | None

Quantization configure.

prefix: str

The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

fastvideo.models.dits.matrixgame3.model.ReplicatedLinear.weight_loader

weight_loader(param: Parameter, loaded_weight: torch.Tensor) -> None

fastvideo.models.dits.matrixgame3.model.ReplicatedLinear.forward

forward(x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]

fastvideo.models.dits.matrixgame3.model.ReplicatedLinear.extra_repr

extra_repr() -> str

fastvideo.models.dits.matrixgame3.model.ReplicatedLinear.get_shape_mapping

get_shape_mapping() -> dict

Get the mapping from (input_shape, output_shape) to layer types.

fastvideo.models.dits.matrixgame3.model.ReplicatedLinear.reset_shape_tracking

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

fastvideo.models.dits.matrixgame3.model.ReplicatedLinear.print_shape_summary

print_shape_summary() -> None

Log a summary of all unique shapes and their layer types.

fastvideo.models.dits.matrixgame3.model.MLP

MLP(input_dim: int, mlp_hidden_dim: int, output_dim: int | None = None, bias: bool = True, act_type: str = 'gelu_pytorch_tanh', dtype: torch.dtype | None = None, prefix: str = '', quant_config: QuantizationConfig | None = None)

MLP for DiT blocks, NO gated linear units

fastvideo.models.dits.matrixgame3.model.MLP.forward

forward(x: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.matrixgame3.model.PatchEmbed

PatchEmbed(patch_size = 16, in_chans = 3, embed_dim = 768, norm_layer = None, flatten = True, bias = True, dtype = None, prefix: str = '')

2D Image to Patch Embedding

Image to Patch Embedding using Conv2d

A convolution based approach to patchifying a 2D image w/ embedding projection.

Based on the impl in https://github.com/google-research/vision_transformer

Hacked together by / Copyright 2020 Ross Wightman

Remove the _assert function in forward function to be compatible with multi-resolution images.

fastvideo.models.dits.matrixgame3.model.PatchEmbed.forward

forward(x)

fastvideo.models.dits.matrixgame3.model.TimestepEmbedder

TimestepEmbedder(hidden_size, act_layer = 'silu', frequency_embedding_size = 256, max_period = 10000, dtype = None, freq_dtype = torch.float32, prefix: str = '')

Embeds scalar timesteps into vector representations.

fastvideo.models.dits.matrixgame3.model.TimestepEmbedder.forward

forward(t: torch.Tensor, timestep_seq_len: int | None = None) -> torch.Tensor

fastvideo.models.dits.matrixgame3.model.ModulateProjection

ModulateProjection(hidden_size: int, factor: int = 2, act_layer: str = 'silu', dtype: torch.dtype | None = None, prefix: str = '')

Modulation layer for DiT blocks.

fastvideo.models.dits.matrixgame3.model.ModulateProjection.forward

forward(x: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.matrixgame3.model.WanCamControlPatchEmbedding

WanCamControlPatchEmbedding(patch_size = (1, 2, 2), in_chans = 384, embed_dim = 2048, bias = True, dtype = None, prefix: str = '')

Lingbot World Patch embedding for Plucker features.

fastvideo.models.dits.matrixgame3.model.WanCamControlPatchEmbedding.forward

forward(x: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.matrixgame3.model.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.models.dits.matrixgame3.model.BaseDiT

BaseDiT(config: DiTConfig, hf_config: dict[str, Any], **kwargs)

fastvideo.models.dits.matrixgame3.model.BaseDiT.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance = None, **kwargs) -> torch.Tensor

fastvideo.models.dits.matrixgame3.model.BaseDiT.preprocess_loaded_weight

preprocess_loaded_weight(name: str, tensor: torch.Tensor) -> torch.Tensor

Adapt a checkpoint tensor before sharding; default models are unchanged.

fastvideo.models.dits.matrixgame3.model.BaseDiT.inference_config_overrides

inference_config_overrides() -> dict[str, Any]

Extra config.json keys a deployable checkpoint of this model needs.

fastvideo.models.dits.matrixgame3.model.BaseDiT.inference_export_expected_shape

inference_export_expected_shape(output_key: str, base_shape: tuple[int, ...]) -> tuple[int, ...]

The exported shape of output_key given the base checkpoint's shape.

fastvideo.models.dits.matrixgame3.model.WanSelfAttention

WanSelfAttention(dim: int, num_heads: int, window_size = (-1, -1), qk_norm = True, eps = 1e-06, parallel_attention = False, quant_config: QuantizationConfig | None = None, prefix: str = '')

fastvideo.models.dits.matrixgame3.model.WanSelfAttention.forward

forward(x: torch.Tensor, context: torch.Tensor, context_lens: int)

Parameters

x: Tensor

Shape [B, L, num_heads, C / num_heads]

seq_lens: Tensor

Shape [B]

grid_sizes: Tensor

Shape [B, 3], the second dimension contains (F, H, W)

freqs: Tensor

Rope freqs, shape [1024, C / num_heads / 2]

fastvideo.models.dits.matrixgame3.model.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.models.dits.matrixgame3.model.MatrixGame3ActionModule

MatrixGame3ActionModule(mouse_dim_in: int = 2, keyboard_dim_in: int = 6, hidden_size: int = 128, img_hidden_size: int = 1536, keyboard_hidden_dim: int = 1024, mouse_hidden_dim: int = 1024, vae_time_compression_ratio: int = 4, windows_size: int = 3, heads_num: int = 16, patch_size: list | None = None, qk_norm: bool = True, qkv_bias: bool = False, rope_dim_list: list | None = None, rope_theta = 256, mouse_qk_dim_list: list | None = None, enable_mouse = True, enable_keyboard = True, blocks: list | None = None)

action module from https://arxiv.org/pdf/2501.08325

fastvideo.models.dits.matrixgame3.model.MatrixGame3ActionModule.patchify

patchify(x, patch_size)

x : (N C T H W)

fastvideo.models.dits.matrixgame3.model.MatrixGame3ActionModule.unpatchify

unpatchify(x, t, h, w, patch_size)

x: (N, T, patch_size**2 * C) imgs: (N, H, W, C)

fastvideo.models.dits.matrixgame3.model.MatrixGame3ActionModule.get_rotary_pos_embed

get_rotary_pos_embed(video_length, height, width, head_dim, rope_dim_list = None)

fastvideo.models.dits.matrixgame3.model.MatrixGame3ActionModule.forward

forward(x: torch.Tensor, tt: int, th: int, tw: int, mouse_condition: torch.Tensor | None = None, keyboard_condition: torch.Tensor | None = None, mouse_cond_memory: torch.Tensor | None = None, keyboard_cond_memory: torch.Tensor | None = None)

hidden_states: B, ttthtw, C mouse_condition: B, N_frames, C1 keyboard_condition: B, N_frames, C2

fastvideo.models.dits.matrixgame3.model.MatrixGame3TimeImageEmbedding

MatrixGame3TimeImageEmbedding(dim: int, time_freq_dim: int)
Source
class MatrixGame3TimeImageEmbedding(nn.Module):

    def __init__(
        self,
        dim: int,
        time_freq_dim: int,
    ):
        super().__init__()

        self.time_embedder = TimestepEmbedder(dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
        self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu")

    def forward(
        self,
        timestep: torch.Tensor,
        encoder_hidden_states: torch.Tensor | None,
        timestep_seq_len: int | None = None,
    ):
        temb = self.time_embedder(timestep, timestep_seq_len)
        timestep_proj = self.time_modulation(temb)

        if encoder_hidden_states is None:
            batch_size = temb.shape[0] if temb.dim() > 1 else timestep.shape[0]
            encoder_hidden_states = torch.zeros(
                (batch_size, 0, temb.shape[-1]),
                device=temb.device,
                dtype=temb.dtype,
            )

        return temb, timestep_proj, encoder_hidden_states

fastvideo.models.dits.matrixgame3.model.MatrixGame3TimeImageEmbedding.forward

forward(timestep: torch.Tensor, encoder_hidden_states: torch.Tensor | None, timestep_seq_len: int | None = None)
Source
def forward(
    self,
    timestep: torch.Tensor,
    encoder_hidden_states: torch.Tensor | None,
    timestep_seq_len: int | None = None,
):
    temb = self.time_embedder(timestep, timestep_seq_len)
    timestep_proj = self.time_modulation(temb)

    if encoder_hidden_states is None:
        batch_size = temb.shape[0] if temb.dim() > 1 else timestep.shape[0]
        encoder_hidden_states = torch.zeros(
            (batch_size, 0, temb.shape[-1]),
            device=temb.device,
            dtype=temb.dtype,
        )

    return temb, timestep_proj, encoder_hidden_states

fastvideo.models.dits.matrixgame3.model.MatrixGame3CrossAttention

MatrixGame3CrossAttention(dim: int, num_heads: int, window_size = (-1, -1), qk_norm = True, eps = 1e-06, parallel_attention = False, quant_config: QuantizationConfig | None = None, prefix: str = '')
Source
class MatrixGame3CrossAttention(WanSelfAttention):

    def forward(self, x, context, context_lens=None):
        r"""
        Args:
            x(Tensor): Shape [B, L1, C]
            context(Tensor): Shape [B, L2, C] - typically 257 image tokens
            context_lens(Tensor): Shape [B]
        """
        b, n, d = x.size(0), self.num_heads, self.head_dim

        q_input = x.to(self.to_q.weight.dtype)
        q = self.norm_q(self.to_q(q_input)[0]).view(b, -1, n, d)

        context_input = context.to(self.to_k.weight.dtype)
        k = self.norm_k(self.to_k(context_input)[0]).view(b, -1, n, d)
        v = self.to_v(context_input)[0].view(b, -1, n, d)

        x = self.attn(q, k, v)

        x = x.flatten(2)
        x, _ = self.to_out(x)
        return x

fastvideo.models.dits.matrixgame3.model.MatrixGame3CrossAttention.forward

forward(x, context, context_lens = None)

Parameters

x: Tensor

Shape [B, L1, C]

context: Tensor

Shape [B, L2, C] - typically 257 image tokens

context_lens: Tensor

Shape [B]

Source
def forward(self, x, context, context_lens=None):
    r"""
    Args:
        x(Tensor): Shape [B, L1, C]
        context(Tensor): Shape [B, L2, C] - typically 257 image tokens
        context_lens(Tensor): Shape [B]
    """
    b, n, d = x.size(0), self.num_heads, self.head_dim

    q_input = x.to(self.to_q.weight.dtype)
    q = self.norm_q(self.to_q(q_input)[0]).view(b, -1, n, d)

    context_input = context.to(self.to_k.weight.dtype)
    k = self.norm_k(self.to_k(context_input)[0]).view(b, -1, n, d)
    v = self.to_v(context_input)[0].view(b, -1, n, d)

    x = self.attn(q, k, v)

    x = x.flatten(2)
    x, _ = self.to_out(x)
    return x

fastvideo.models.dits.matrixgame3.model.MatrixGame3TransformerBlock

MatrixGame3TransformerBlock(dim: int, ffn_dim: int, num_heads: int, qk_norm: str = 'rms_norm_across_heads', cross_attn_norm: bool = False, eps: float = 1e-06, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '', action_config: dict | None = None, block_id: int | None = None, use_memory: bool = False)
Source
class MatrixGame3TransformerBlock(nn.Module):

    def __init__(
        self,
        dim: int,
        ffn_dim: int,
        num_heads: int,
        qk_norm: str = "rms_norm_across_heads",
        cross_attn_norm: bool = False,
        eps: float = 1e-6,
        supported_attention_backends: tuple[AttentionBackendEnum, ...]
        | None = None,
        prefix: str = "",
        action_config: dict | None = None,
        block_id: int | None = None,
        use_memory: bool = False,
    ):
        super().__init__()
        action_config = action_config or {}
        self.use_memory = use_memory

        # 1. Self-attention
        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
        self.to_q = ReplicatedLinear(dim, dim, bias=True)
        self.to_k = ReplicatedLinear(dim, dim, bias=True)
        self.to_v = ReplicatedLinear(dim, dim, bias=True)

        self.to_out = ReplicatedLinear(dim, dim, bias=True)
        self.attn1 = DistributedAttention(
            num_heads=num_heads,
            head_size=dim // num_heads,
            causal=False,
            supported_attention_backends=supported_attention_backends,
            prefix=f"{prefix}.attn1",
        )
        self.hidden_dim = dim
        self.num_attention_heads = num_heads
        dim_head = dim // num_heads
        if qk_norm == "rms_norm":
            self.norm_q = RMSNorm(dim_head, eps=eps)
            self.norm_k = RMSNorm(dim_head, eps=eps)
        elif qk_norm == "rms_norm_across_heads":
            # LTX applies qk norm across all heads
            self.norm_q = RMSNorm(dim, eps=eps)
            self.norm_k = RMSNorm(dim, eps=eps)
        else:
            print("QK Norm type not supported")
            raise Exception
        assert cross_attn_norm is True
        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim,
            norm_type="layer",
            eps=eps,
            elementwise_affine=True,
            dtype=torch.float32,
            compute_dtype=torch.float32,
        )

        # 2. Cross-attention
        self.attn2 = MatrixGame3CrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim,
            norm_type="layer",
            eps=eps,
            elementwise_affine=False,
            dtype=torch.float32,
            compute_dtype=torch.float32,
        )

        # 2.1. Action Module Integration
        enabled_blocks = set(action_config.get("blocks", [])) if action_config else set()
        self.use_action_module = len(action_config) > 0 and (block_id is None or block_id in enabled_blocks)
        if self.use_action_module:
            self.action_model = MatrixGame3ActionModule(**action_config)
        else:
            self.action_model = None

        # 3. Feed-forward
        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
        self.mlp_residual = ScaleResidual()

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
        if self.use_memory:
            self.cam_injector_layer1 = nn.Linear(dim, dim)
            self.cam_injector_layer2 = nn.Linear(dim, dim)
            self.cam_scale_layer = nn.Linear(dim, dim)
            self.cam_shift_layer = nn.Linear(dim, dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        freqs: torch.Tensor,
        # Action Module specific args
        grid_sizes: tuple[int, int, int],
        mouse_cond: torch.Tensor | None = None,
        keyboard_cond: torch.Tensor | None = None,
        mouse_cond_memory: torch.Tensor | None = None,
        keyboard_cond_memory: torch.Tensor | None = None,
        plucker_emb: torch.Tensor | None = None,
        memory_length: int = 0,
        memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None,
        predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None,
    ) -> torch.Tensor:
        if hidden_states.dim() == 4:
            hidden_states = hidden_states.squeeze(1)
        orig_dtype = hidden_states.dtype

        if temb.dim() == 4:
            # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
            (
                shift_msa,
                scale_msa,
                gate_msa,
                c_shift_msa,
                c_scale_msa,
                c_gate_msa,
            ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)
            # batch_size, seq_len, 1, inner_dim
            shift_msa = shift_msa.squeeze(2)
            scale_msa = scale_msa.squeeze(2)
            gate_msa = gate_msa.squeeze(2)
            c_shift_msa = c_shift_msa.squeeze(2)
            c_scale_msa = c_scale_msa.squeeze(2)
            c_gate_msa = c_gate_msa.squeeze(2)
        else:
            # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
            e = self.scale_shift_table + temb.float()
            (
                shift_msa,
                scale_msa,
                gate_msa,
                c_shift_msa,
                c_scale_msa,
                c_gate_msa,
            ) = e.chunk(6, dim=1)
        assert shift_msa.dtype == torch.float32

        # 1. Self-attention
        norm_hidden_states = (self.norm1(hidden_states).float() * (1 + scale_msa) + shift_msa).to(orig_dtype)
        query, _ = self.to_q(norm_hidden_states)
        key, _ = self.to_k(norm_hidden_states)
        value, _ = self.to_v(norm_hidden_states)

        if self.norm_q is not None:
            query = self.norm_q(query)
        if self.norm_k is not None:
            key = self.norm_k(key)

        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))

        # Apply rotary embeddings
        grid_frames, grid_height, grid_width = grid_sizes

        if memory_length > 0:
            hw = grid_height * grid_width
            query_memory = query[:, :memory_length * hw]
            key_memory = key[:, :memory_length * hw]
            query_pred = query[:, memory_length * hw:]
            key_pred = key[:, memory_length * hw:]

            # Indices are pre-built device tensors by MatrixGame3WanModel.forward.
            mem_indices = memory_latent_idx
            if mem_indices is None:
                mem_indices = torch.arange(memory_length, device=query.device, dtype=torch.long)
            pred_indices = predict_latent_idx
            if pred_indices is None:
                pred_indices = torch.arange(grid_frames - memory_length, device=query.device, dtype=torch.long)

            query_memory = _apply_rope_with_frame_indices(query_memory,
                                                          freqs,
                                                          height=grid_height,
                                                          width=grid_width,
                                                          frame_indices=mem_indices)
            key_memory = _apply_rope_with_frame_indices(key_memory,
                                                        freqs,
                                                        height=grid_height,
                                                        width=grid_width,
                                                        frame_indices=mem_indices)
            query_pred = _apply_rope_with_frame_indices(query_pred,
                                                        freqs,
                                                        height=grid_height,
                                                        width=grid_width,
                                                        frame_indices=pred_indices)
            key_pred = _apply_rope_with_frame_indices(key_pred,
                                                      freqs,
                                                      height=grid_height,
                                                      width=grid_width,
                                                      frame_indices=pred_indices)
            query = torch.cat([query_memory, query_pred], dim=1)
            key = torch.cat([key_memory, key_pred], dim=1)
        else:
            pred_indices = predict_latent_idx
            if pred_indices is None:
                pred_indices = torch.arange(grid_frames, device=query.device, dtype=torch.long)
            query = _apply_rope_with_frame_indices(query,
                                                   freqs,
                                                   height=grid_height,
                                                   width=grid_width,
                                                   frame_indices=pred_indices)
            key = _apply_rope_with_frame_indices(key,
                                                 freqs,
                                                 height=grid_height,
                                                 width=grid_width,
                                                 frame_indices=pred_indices)

        attn_output, _ = self.attn1(query, key, value)
        attn_output = attn_output.flatten(2)
        attn_output = attn_output.to(self.to_out.weight.dtype)
        attn_output, _ = self.to_out(attn_output)
        attn_output = attn_output.squeeze(1)

        hidden_states = hidden_states.float() + attn_output.float() * gate_msa.float()
        if self.use_memory and plucker_emb is not None:
            c2ws_hidden_states = self.cam_injector_layer2(
                torch.nn.functional.silu(self.cam_injector_layer1(plucker_emb)))
            c2ws_hidden_states = c2ws_hidden_states + plucker_emb
            cam_scale = self.cam_scale_layer(c2ws_hidden_states)
            cam_shift = self.cam_shift_layer(c2ws_hidden_states)
            hidden_states = (1.0 + cam_scale) * hidden_states + cam_shift
        norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states)

        # 2. Cross-attention
        attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)

        # - use_memory / action path: x = norm3(x); x = x + cross_attn(x)
        # - plain path:              x = x + cross_attn(norm3(x))
        cross_attn_residual_base = (norm_hidden_states if
                                    (mouse_cond is not None or self.use_memory) else hidden_states)
        hidden_states = cross_attn_residual_base + attn_output
        norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float())
        norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa

        # ================= Action Module =================
        if self.action_model is not None and (mouse_cond is not None or keyboard_cond is not None):
            action_dtype = self.ffn.fc_in.weight.dtype
            # grid_sizes is expected to be [F, H, W]
            hidden_states = self.action_model(
                hidden_states.to(action_dtype),
                grid_sizes[0],
                grid_sizes[1],
                grid_sizes[2],
                mouse_cond,
                keyboard_cond,
                mouse_cond_memory=mouse_cond_memory,
                keyboard_cond_memory=keyboard_cond_memory,
            )

            norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float())
            norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa
        # =================================================

        # 3. Feed-forward
        ff_output = self.ffn(norm_hidden_states.to(self.ffn.fc_in.weight.dtype))
        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
        hidden_states = hidden_states.to(orig_dtype)

        return hidden_states

fastvideo.models.dits.matrixgame3.model.MatrixGame3TransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs: torch.Tensor, grid_sizes: tuple[int, int, int], mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, mouse_cond_memory: torch.Tensor | None = None, keyboard_cond_memory: torch.Tensor | None = None, plucker_emb: torch.Tensor | None = None, memory_length: int = 0, memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None, predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    temb: torch.Tensor,
    freqs: torch.Tensor,
    # Action Module specific args
    grid_sizes: tuple[int, int, int],
    mouse_cond: torch.Tensor | None = None,
    keyboard_cond: torch.Tensor | None = None,
    mouse_cond_memory: torch.Tensor | None = None,
    keyboard_cond_memory: torch.Tensor | None = None,
    plucker_emb: torch.Tensor | None = None,
    memory_length: int = 0,
    memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None,
    predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None,
) -> torch.Tensor:
    if hidden_states.dim() == 4:
        hidden_states = hidden_states.squeeze(1)
    orig_dtype = hidden_states.dtype

    if temb.dim() == 4:
        # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
        (
            shift_msa,
            scale_msa,
            gate_msa,
            c_shift_msa,
            c_scale_msa,
            c_gate_msa,
        ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2)
        # batch_size, seq_len, 1, inner_dim
        shift_msa = shift_msa.squeeze(2)
        scale_msa = scale_msa.squeeze(2)
        gate_msa = gate_msa.squeeze(2)
        c_shift_msa = c_shift_msa.squeeze(2)
        c_scale_msa = c_scale_msa.squeeze(2)
        c_gate_msa = c_gate_msa.squeeze(2)
    else:
        # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
        e = self.scale_shift_table + temb.float()
        (
            shift_msa,
            scale_msa,
            gate_msa,
            c_shift_msa,
            c_scale_msa,
            c_gate_msa,
        ) = e.chunk(6, dim=1)
    assert shift_msa.dtype == torch.float32

    # 1. Self-attention
    norm_hidden_states = (self.norm1(hidden_states).float() * (1 + scale_msa) + shift_msa).to(orig_dtype)
    query, _ = self.to_q(norm_hidden_states)
    key, _ = self.to_k(norm_hidden_states)
    value, _ = self.to_v(norm_hidden_states)

    if self.norm_q is not None:
        query = self.norm_q(query)
    if self.norm_k is not None:
        key = self.norm_k(key)

    query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
    key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
    value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))

    # Apply rotary embeddings
    grid_frames, grid_height, grid_width = grid_sizes

    if memory_length > 0:
        hw = grid_height * grid_width
        query_memory = query[:, :memory_length * hw]
        key_memory = key[:, :memory_length * hw]
        query_pred = query[:, memory_length * hw:]
        key_pred = key[:, memory_length * hw:]

        # Indices are pre-built device tensors by MatrixGame3WanModel.forward.
        mem_indices = memory_latent_idx
        if mem_indices is None:
            mem_indices = torch.arange(memory_length, device=query.device, dtype=torch.long)
        pred_indices = predict_latent_idx
        if pred_indices is None:
            pred_indices = torch.arange(grid_frames - memory_length, device=query.device, dtype=torch.long)

        query_memory = _apply_rope_with_frame_indices(query_memory,
                                                      freqs,
                                                      height=grid_height,
                                                      width=grid_width,
                                                      frame_indices=mem_indices)
        key_memory = _apply_rope_with_frame_indices(key_memory,
                                                    freqs,
                                                    height=grid_height,
                                                    width=grid_width,
                                                    frame_indices=mem_indices)
        query_pred = _apply_rope_with_frame_indices(query_pred,
                                                    freqs,
                                                    height=grid_height,
                                                    width=grid_width,
                                                    frame_indices=pred_indices)
        key_pred = _apply_rope_with_frame_indices(key_pred,
                                                  freqs,
                                                  height=grid_height,
                                                  width=grid_width,
                                                  frame_indices=pred_indices)
        query = torch.cat([query_memory, query_pred], dim=1)
        key = torch.cat([key_memory, key_pred], dim=1)
    else:
        pred_indices = predict_latent_idx
        if pred_indices is None:
            pred_indices = torch.arange(grid_frames, device=query.device, dtype=torch.long)
        query = _apply_rope_with_frame_indices(query,
                                               freqs,
                                               height=grid_height,
                                               width=grid_width,
                                               frame_indices=pred_indices)
        key = _apply_rope_with_frame_indices(key,
                                             freqs,
                                             height=grid_height,
                                             width=grid_width,
                                             frame_indices=pred_indices)

    attn_output, _ = self.attn1(query, key, value)
    attn_output = attn_output.flatten(2)
    attn_output = attn_output.to(self.to_out.weight.dtype)
    attn_output, _ = self.to_out(attn_output)
    attn_output = attn_output.squeeze(1)

    hidden_states = hidden_states.float() + attn_output.float() * gate_msa.float()
    if self.use_memory and plucker_emb is not None:
        c2ws_hidden_states = self.cam_injector_layer2(
            torch.nn.functional.silu(self.cam_injector_layer1(plucker_emb)))
        c2ws_hidden_states = c2ws_hidden_states + plucker_emb
        cam_scale = self.cam_scale_layer(c2ws_hidden_states)
        cam_shift = self.cam_shift_layer(c2ws_hidden_states)
        hidden_states = (1.0 + cam_scale) * hidden_states + cam_shift
    norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states)

    # 2. Cross-attention
    attn_output = self.attn2(norm_hidden_states, context=encoder_hidden_states, context_lens=None)

    # - use_memory / action path: x = norm3(x); x = x + cross_attn(x)
    # - plain path:              x = x + cross_attn(norm3(x))
    cross_attn_residual_base = (norm_hidden_states if
                                (mouse_cond is not None or self.use_memory) else hidden_states)
    hidden_states = cross_attn_residual_base + attn_output
    norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float())
    norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa

    # ================= Action Module =================
    if self.action_model is not None and (mouse_cond is not None or keyboard_cond is not None):
        action_dtype = self.ffn.fc_in.weight.dtype
        # grid_sizes is expected to be [F, H, W]
        hidden_states = self.action_model(
            hidden_states.to(action_dtype),
            grid_sizes[0],
            grid_sizes[1],
            grid_sizes[2],
            mouse_cond,
            keyboard_cond,
            mouse_cond_memory=mouse_cond_memory,
            keyboard_cond_memory=keyboard_cond_memory,
        )

        norm_hidden_states = self.cross_attn_residual_norm.norm(hidden_states.float())
        norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa
    # =================================================

    # 3. Feed-forward
    ff_output = self.ffn(norm_hidden_states.to(self.ffn.fc_in.weight.dtype))
    hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
    hidden_states = hidden_states.to(orig_dtype)

    return hidden_states

fastvideo.models.dits.matrixgame3.model.MatrixGame3WanModel

MatrixGame3WanModel(config: MatrixGame3WanVideoConfig, hf_config: dict[str, Any], **kwargs)
Source
class MatrixGame3WanModel(BaseDiT):
    # Marker for action input support (Matrix-Game)
    supports_action_input = True

    _fsdp_shard_conditions = _DEFAULT_MATRIXGAME3_CONFIG._fsdp_shard_conditions
    _compile_conditions = _DEFAULT_MATRIXGAME3_CONFIG._compile_conditions
    _supported_attention_backends = (_DEFAULT_MATRIXGAME3_CONFIG._supported_attention_backends)
    param_names_mapping = _DEFAULT_MATRIXGAME3_CONFIG.param_names_mapping
    reverse_param_names_mapping = (_DEFAULT_MATRIXGAME3_CONFIG.reverse_param_names_mapping)
    lora_param_names_mapping = (_DEFAULT_MATRIXGAME3_CONFIG.lora_param_names_mapping)

    def __init__(
        self,
        config: MatrixGame3WanVideoConfig,
        hf_config: dict[str, Any],
        **kwargs,
    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        inner_dim = config.num_attention_heads * config.attention_head_dim
        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.num_channels_latents = config.out_channels
        self.patch_size = config.patch_size

        # 1. Patch & position embedding
        self.patch_embedding = PatchEmbed(
            in_chans=config.in_channels,
            embed_dim=inner_dim,
            patch_size=config.patch_size,
            flatten=False,
        )

        # 2. Condition embeddings
        self.condition_embedder = MatrixGame3TimeImageEmbedding(
            dim=inner_dim,
            time_freq_dim=config.freq_dim,
        )
        self.text_embedding = nn.Sequential(
            nn.Linear(config.text_dim, inner_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(inner_dim, inner_dim),
        )
        self.use_memory = getattr(config, "use_memory", False)
        self.sigma_theta = float(getattr(config, "sigma_theta", 0.0))
        self.camera_patch_embedding = None
        self.c2ws_hidden_states_layer1 = None
        self.c2ws_hidden_states_layer2 = None
        if self.use_memory:
            camera_in_channels = getattr(config, "camera_embed_in_channels", 1536)
            self.camera_patch_embedding = WanCamControlPatchEmbedding(
                patch_size=config.patch_size,
                in_chans=camera_in_channels,
                embed_dim=inner_dim,
            )
            self.c2ws_hidden_states_layer1 = nn.Linear(inner_dim, inner_dim)
            self.c2ws_hidden_states_layer2 = nn.Linear(inner_dim, inner_dim)

        # 2.1. Get action config
        self.action_config = getattr(config, "action_config", {})

        # 3. Transformer blocks
        self.blocks = nn.ModuleList([
            MatrixGame3TransformerBlock(
                inner_dim,
                config.ffn_dim,
                config.num_attention_heads,
                config.qk_norm,
                config.cross_attn_norm,
                config.eps,
                self._supported_attention_backends,
                prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}",
                action_config=self.action_config,
                block_id=i,
                use_memory=self.use_memory,
            ) for i in range(config.num_layers)
        ])

        # 4. Output norm & projection
        self.norm_out = LayerNormScaleShift(
            inner_dim,
            norm_type="layer",
            eps=config.eps,
            elementwise_affine=False,
            dtype=torch.float32,
            compute_dtype=torch.float32,
        )
        self.proj_out = nn.Linear(inner_dim, config.out_channels * math.prod(config.patch_size))
        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)

        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
        timestep: torch.LongTensor,
        # Action inputs
        mouse_cond: torch.Tensor | None = None,
        keyboard_cond: torch.Tensor | None = None,
        x_memory: torch.Tensor | None = None,
        timestep_memory: torch.Tensor | None = None,
        mouse_cond_memory: torch.Tensor | None = None,
        keyboard_cond_memory: torch.Tensor | None = None,
        c2ws_plucker_emb: torch.Tensor | None = None,
        memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None,
        predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]

        memory_length = 0
        if x_memory is not None:
            memory_length = x_memory.shape[2]
            hidden_states = torch.cat([x_memory.to(hidden_states.dtype), hidden_states], dim=2)

        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        p_t, p_h, p_w = self.patch_size
        post_patch_num_frames = num_frames // p_t
        post_patch_height = height // p_h
        post_patch_width = width // p_w
        max_frame_index = post_patch_num_frames
        if isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2:
            max_frame_index = max(max_frame_index, int(predict_latent_idx[1]))
        elif predict_latent_idx is not None:
            max_frame_index = max(max_frame_index, int(max(predict_latent_idx)) + 1)
        if memory_latent_idx is not None and len(memory_latent_idx) > 0:
            max_frame_index = max(max_frame_index, int(max(memory_latent_idx)) + 1)

        # Pre-build RoPE frame index tensors once per forward so blocks don't
        # do `torch.tensor(list, device=cuda)` per call (host sync).
        device_idx = hidden_states.device
        if predict_latent_idx is None:
            predict_latent_idx = None
        elif torch.is_tensor(predict_latent_idx):
            predict_latent_idx = predict_latent_idx.to(device=device_idx, dtype=torch.long)
        elif isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2:
            predict_latent_idx = torch.arange(
                predict_latent_idx[0],
                predict_latent_idx[1],
                device=device_idx,
                dtype=torch.long,
            )
        else:
            predict_latent_idx = torch.tensor(
                list(predict_latent_idx),
                device=device_idx,
                dtype=torch.long,
            )

        if memory_latent_idx is None or (not torch.is_tensor(memory_latent_idx) and len(memory_latent_idx) == 0):
            memory_latent_idx = None
        elif torch.is_tensor(memory_latent_idx):
            memory_latent_idx = memory_latent_idx.to(device=device_idx, dtype=torch.long)
        else:
            memory_latent_idx = torch.tensor(
                list(memory_latent_idx),
                device=device_idx,
                dtype=torch.long,
            )

        head_dim = self.hidden_size // self.num_attention_heads
        freqs = _build_rope_freqs(
            max_seq_len=(2048 if self.use_memory else 1024),
            head_dim=head_dim,
            num_heads=self.num_attention_heads,
            sigma_theta=(self.sigma_theta if self.use_memory else 0.0),
            device=hidden_states.device,
        )
        max_seq_len = freqs.shape[1] if freqs.dim() == 3 else freqs.shape[0]
        if max_frame_index > max_seq_len or post_patch_height > max_seq_len or post_patch_width > max_seq_len:
            raise ValueError(f"got frames={max_frame_index}, height={post_patch_height}, width={post_patch_width}, "
                             f"max_seq_len={max_seq_len}")

        hidden_states = self.patch_embedding(hidden_states)
        hidden_states = hidden_states.flatten(2).transpose(1, 2)

        plucker_emb = None
        if c2ws_plucker_emb is not None and self.camera_patch_embedding is not None:
            if memory_length > 0 and c2ws_plucker_emb.shape[2] == post_patch_num_frames - memory_length:
                zeros = torch.zeros(c2ws_plucker_emb.shape[0],
                                    c2ws_plucker_emb.shape[1],
                                    memory_length,
                                    c2ws_plucker_emb.shape[3],
                                    c2ws_plucker_emb.shape[4],
                                    device=c2ws_plucker_emb.device,
                                    dtype=c2ws_plucker_emb.dtype)
                c2ws_plucker_emb = torch.cat([zeros, c2ws_plucker_emb], dim=2)
            target_frames = post_patch_num_frames * p_t
            target_height = post_patch_height * p_h
            target_width = post_patch_width * p_w
            if c2ws_plucker_emb.shape[2] > target_frames or c2ws_plucker_emb.shape[
                    3] > target_height or c2ws_plucker_emb.shape[4] > target_width:
                c2ws_plucker_emb = c2ws_plucker_emb[:, :, :target_frames, :target_height, :target_width]
            if c2ws_plucker_emb.shape[2:] != (target_frames, target_height, target_width):
                raise ValueError("Camera conditioning must align with the patchified latent grid, "
                                 f"got camera shape {tuple(c2ws_plucker_emb.shape[2:])} vs target "
                                 f"{(target_frames, target_height, target_width)}")
            plucker_emb = self.camera_patch_embedding(c2ws_plucker_emb.to(hidden_states.dtype))
            plucker_hidden = self.c2ws_hidden_states_layer2(
                torch.nn.functional.silu(self.c2ws_hidden_states_layer1(plucker_emb)))
            plucker_emb = plucker_emb + plucker_hidden

        timestep_tokens = timestep
        if timestep_tokens.dim() == 0:
            timestep_tokens = timestep_tokens.unsqueeze(0)
        if timestep_tokens.dim() == 1:
            timestep_tokens = timestep_tokens.unsqueeze(1).repeat(
                1, post_patch_num_frames * post_patch_height * post_patch_width)
        elif timestep_tokens.dim() == 2 and timestep_tokens.shape[1] == num_frames:
            timestep_tokens = (timestep_tokens[:, :, None,
                                               None].repeat(1, 1, post_patch_height,
                                                            post_patch_width).reshape(timestep_tokens.shape[0], -1))
        if memory_length > 0:
            if timestep_memory is None:
                raise ValueError("MatrixGame3 requires timestep_memory when x_memory is provided")
            timestep_tokens = torch.cat([timestep_memory.to(timestep_tokens.dtype), timestep_tokens], dim=1)

        temb, timestep_proj, encoder_hidden_states = self.condition_embedder(
            timestep_tokens.flatten(),
            encoder_hidden_states,
            timestep_seq_len=timestep_tokens.shape[1],
        )
        if timestep_proj.dim() == 3 and timestep_proj.shape[-1] % 6 == 0:
            timestep_proj = timestep_proj.unflatten(2, (6, -1))

        if encoder_hidden_states is not None:
            if isinstance(encoder_hidden_states, list):
                encoder_hidden_states = encoder_hidden_states[0]
            elif encoder_hidden_states.ndim == 2:
                encoder_hidden_states = encoder_hidden_states.unsqueeze(0)
            if (encoder_hidden_states.ndim == 3
                    and encoder_hidden_states.shape[-1] == self.text_embedding[0].in_features):
                encoder_hidden_states = self.text_embedding(encoder_hidden_states.to(hidden_states.dtype))

        # [F, H, W] for the ActionModule
        grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)

        # Blocks
        for block in self.blocks:
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                hidden_states = self._gradient_checkpointing_func(
                    block,
                    hidden_states,
                    encoder_hidden_states,
                    timestep_proj,
                    freqs,
                    grid_sizes=grid_sizes,
                    mouse_cond=mouse_cond,
                    keyboard_cond=keyboard_cond,
                    mouse_cond_memory=mouse_cond_memory,
                    keyboard_cond_memory=keyboard_cond_memory,
                    plucker_emb=plucker_emb,
                    memory_length=memory_length,
                    memory_latent_idx=memory_latent_idx,
                    predict_latent_idx=predict_latent_idx,
                )
            else:
                hidden_states = block(
                    hidden_states,
                    encoder_hidden_states,
                    timestep_proj,
                    freqs,
                    grid_sizes=grid_sizes,
                    mouse_cond=mouse_cond,
                    keyboard_cond=keyboard_cond,
                    mouse_cond_memory=mouse_cond_memory,
                    keyboard_cond_memory=keyboard_cond_memory,
                    plucker_emb=plucker_emb,
                    memory_length=memory_length,
                    memory_latent_idx=memory_latent_idx,
                    predict_latent_idx=predict_latent_idx,
                )

        # Output
        if temb.dim() == 3:
            shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
            shift = shift.squeeze(2)
            scale = scale.squeeze(2)
        else:
            shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
        hidden_states = self.norm_out(hidden_states, shift, scale)
        hidden_states = self.proj_out(hidden_states)

        hidden_states = hidden_states.reshape(
            batch_size,
            post_patch_num_frames,
            post_patch_height,
            post_patch_width,
            p_t,
            p_h,
            p_w,
            -1,
        )
        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

        if memory_length > 0:
            output = output[:, :, memory_length:]

        return output

fastvideo.models.dits.matrixgame3.model.MatrixGame3WanModel.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, x_memory: torch.Tensor | None = None, timestep_memory: torch.Tensor | None = None, mouse_cond_memory: torch.Tensor | None = None, keyboard_cond_memory: torch.Tensor | None = None, c2ws_plucker_emb: torch.Tensor | None = None, memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None, predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None, **kwargs) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor | list[torch.Tensor],
    timestep: torch.LongTensor,
    # Action inputs
    mouse_cond: torch.Tensor | None = None,
    keyboard_cond: torch.Tensor | None = None,
    x_memory: torch.Tensor | None = None,
    timestep_memory: torch.Tensor | None = None,
    mouse_cond_memory: torch.Tensor | None = None,
    keyboard_cond_memory: torch.Tensor | None = None,
    c2ws_plucker_emb: torch.Tensor | None = None,
    memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None,
    predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None,
    **kwargs,
) -> torch.Tensor:
    if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states[0]

    memory_length = 0
    if x_memory is not None:
        memory_length = x_memory.shape[2]
        hidden_states = torch.cat([x_memory.to(hidden_states.dtype), hidden_states], dim=2)

    batch_size, num_channels, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.patch_size
    post_patch_num_frames = num_frames // p_t
    post_patch_height = height // p_h
    post_patch_width = width // p_w
    max_frame_index = post_patch_num_frames
    if isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2:
        max_frame_index = max(max_frame_index, int(predict_latent_idx[1]))
    elif predict_latent_idx is not None:
        max_frame_index = max(max_frame_index, int(max(predict_latent_idx)) + 1)
    if memory_latent_idx is not None and len(memory_latent_idx) > 0:
        max_frame_index = max(max_frame_index, int(max(memory_latent_idx)) + 1)

    # Pre-build RoPE frame index tensors once per forward so blocks don't
    # do `torch.tensor(list, device=cuda)` per call (host sync).
    device_idx = hidden_states.device
    if predict_latent_idx is None:
        predict_latent_idx = None
    elif torch.is_tensor(predict_latent_idx):
        predict_latent_idx = predict_latent_idx.to(device=device_idx, dtype=torch.long)
    elif isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2:
        predict_latent_idx = torch.arange(
            predict_latent_idx[0],
            predict_latent_idx[1],
            device=device_idx,
            dtype=torch.long,
        )
    else:
        predict_latent_idx = torch.tensor(
            list(predict_latent_idx),
            device=device_idx,
            dtype=torch.long,
        )

    if memory_latent_idx is None or (not torch.is_tensor(memory_latent_idx) and len(memory_latent_idx) == 0):
        memory_latent_idx = None
    elif torch.is_tensor(memory_latent_idx):
        memory_latent_idx = memory_latent_idx.to(device=device_idx, dtype=torch.long)
    else:
        memory_latent_idx = torch.tensor(
            list(memory_latent_idx),
            device=device_idx,
            dtype=torch.long,
        )

    head_dim = self.hidden_size // self.num_attention_heads
    freqs = _build_rope_freqs(
        max_seq_len=(2048 if self.use_memory else 1024),
        head_dim=head_dim,
        num_heads=self.num_attention_heads,
        sigma_theta=(self.sigma_theta if self.use_memory else 0.0),
        device=hidden_states.device,
    )
    max_seq_len = freqs.shape[1] if freqs.dim() == 3 else freqs.shape[0]
    if max_frame_index > max_seq_len or post_patch_height > max_seq_len or post_patch_width > max_seq_len:
        raise ValueError(f"got frames={max_frame_index}, height={post_patch_height}, width={post_patch_width}, "
                         f"max_seq_len={max_seq_len}")

    hidden_states = self.patch_embedding(hidden_states)
    hidden_states = hidden_states.flatten(2).transpose(1, 2)

    plucker_emb = None
    if c2ws_plucker_emb is not None and self.camera_patch_embedding is not None:
        if memory_length > 0 and c2ws_plucker_emb.shape[2] == post_patch_num_frames - memory_length:
            zeros = torch.zeros(c2ws_plucker_emb.shape[0],
                                c2ws_plucker_emb.shape[1],
                                memory_length,
                                c2ws_plucker_emb.shape[3],
                                c2ws_plucker_emb.shape[4],
                                device=c2ws_plucker_emb.device,
                                dtype=c2ws_plucker_emb.dtype)
            c2ws_plucker_emb = torch.cat([zeros, c2ws_plucker_emb], dim=2)
        target_frames = post_patch_num_frames * p_t
        target_height = post_patch_height * p_h
        target_width = post_patch_width * p_w
        if c2ws_plucker_emb.shape[2] > target_frames or c2ws_plucker_emb.shape[
                3] > target_height or c2ws_plucker_emb.shape[4] > target_width:
            c2ws_plucker_emb = c2ws_plucker_emb[:, :, :target_frames, :target_height, :target_width]
        if c2ws_plucker_emb.shape[2:] != (target_frames, target_height, target_width):
            raise ValueError("Camera conditioning must align with the patchified latent grid, "
                             f"got camera shape {tuple(c2ws_plucker_emb.shape[2:])} vs target "
                             f"{(target_frames, target_height, target_width)}")
        plucker_emb = self.camera_patch_embedding(c2ws_plucker_emb.to(hidden_states.dtype))
        plucker_hidden = self.c2ws_hidden_states_layer2(
            torch.nn.functional.silu(self.c2ws_hidden_states_layer1(plucker_emb)))
        plucker_emb = plucker_emb + plucker_hidden

    timestep_tokens = timestep
    if timestep_tokens.dim() == 0:
        timestep_tokens = timestep_tokens.unsqueeze(0)
    if timestep_tokens.dim() == 1:
        timestep_tokens = timestep_tokens.unsqueeze(1).repeat(
            1, post_patch_num_frames * post_patch_height * post_patch_width)
    elif timestep_tokens.dim() == 2 and timestep_tokens.shape[1] == num_frames:
        timestep_tokens = (timestep_tokens[:, :, None,
                                           None].repeat(1, 1, post_patch_height,
                                                        post_patch_width).reshape(timestep_tokens.shape[0], -1))
    if memory_length > 0:
        if timestep_memory is None:
            raise ValueError("MatrixGame3 requires timestep_memory when x_memory is provided")
        timestep_tokens = torch.cat([timestep_memory.to(timestep_tokens.dtype), timestep_tokens], dim=1)

    temb, timestep_proj, encoder_hidden_states = self.condition_embedder(
        timestep_tokens.flatten(),
        encoder_hidden_states,
        timestep_seq_len=timestep_tokens.shape[1],
    )
    if timestep_proj.dim() == 3 and timestep_proj.shape[-1] % 6 == 0:
        timestep_proj = timestep_proj.unflatten(2, (6, -1))

    if encoder_hidden_states is not None:
        if isinstance(encoder_hidden_states, list):
            encoder_hidden_states = encoder_hidden_states[0]
        elif encoder_hidden_states.ndim == 2:
            encoder_hidden_states = encoder_hidden_states.unsqueeze(0)
        if (encoder_hidden_states.ndim == 3
                and encoder_hidden_states.shape[-1] == self.text_embedding[0].in_features):
            encoder_hidden_states = self.text_embedding(encoder_hidden_states.to(hidden_states.dtype))

    # [F, H, W] for the ActionModule
    grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)

    # Blocks
    for block in self.blocks:
        if torch.is_grad_enabled() and self.gradient_checkpointing:
            hidden_states = self._gradient_checkpointing_func(
                block,
                hidden_states,
                encoder_hidden_states,
                timestep_proj,
                freqs,
                grid_sizes=grid_sizes,
                mouse_cond=mouse_cond,
                keyboard_cond=keyboard_cond,
                mouse_cond_memory=mouse_cond_memory,
                keyboard_cond_memory=keyboard_cond_memory,
                plucker_emb=plucker_emb,
                memory_length=memory_length,
                memory_latent_idx=memory_latent_idx,
                predict_latent_idx=predict_latent_idx,
            )
        else:
            hidden_states = block(
                hidden_states,
                encoder_hidden_states,
                timestep_proj,
                freqs,
                grid_sizes=grid_sizes,
                mouse_cond=mouse_cond,
                keyboard_cond=keyboard_cond,
                mouse_cond_memory=mouse_cond_memory,
                keyboard_cond_memory=keyboard_cond_memory,
                plucker_emb=plucker_emb,
                memory_length=memory_length,
                memory_latent_idx=memory_latent_idx,
                predict_latent_idx=predict_latent_idx,
            )

    # Output
    if temb.dim() == 3:
        shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
        shift = shift.squeeze(2)
        scale = scale.squeeze(2)
    else:
        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
    hidden_states = self.norm_out(hidden_states, shift, scale)
    hidden_states = self.proj_out(hidden_states)

    hidden_states = hidden_states.reshape(
        batch_size,
        post_patch_num_frames,
        post_patch_height,
        post_patch_width,
        p_t,
        p_h,
        p_w,
        -1,
    )
    hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
    output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

    if memory_length > 0:
        output = output[:, :, memory_length:]

    return output