--- title: causal_model ---

fastvideo.models.dits.matrixgame2.causal_model.get_forward_context

get_forward_context() -> ForwardContext

Get the current forward context.

fastvideo.models.dits.matrixgame2.causal_model.LocalAttention

LocalAttention(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, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

fastvideo.models.dits.matrixgame2.causal_model.LocalAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor

Query tensor of shape [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor of shape [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor

torch.Tensor: Output tensor after local attention

fastvideo.models.dits.matrixgame2.causal_model.MatrixGame2WanVideoConfig

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

fastvideo.models.dits.matrixgame2.causal_model.get_sp_world_size

get_sp_world_size() -> int

Return world size for the sequence model parallel group.

fastvideo.models.dits.matrixgame2.causal_model.FP32LayerNorm

FP32LayerNorm()

fastvideo.models.dits.matrixgame2.causal_model.FP32LayerNorm.forward

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

fastvideo.models.dits.matrixgame2.causal_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.matrixgame2.causal_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.matrixgame2.causal_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.matrixgame2.causal_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.matrixgame2.causal_model.RMSNorm.extra_repr

extra_repr() -> str

fastvideo.models.dits.matrixgame2.causal_model.ScaleResidual

ScaleResidual(prefix: str = '')

Applies gated residual connection.

fastvideo.models.dits.matrixgame2.causal_model.ScaleResidual.forward

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

Apply gated residual connection.

fastvideo.models.dits.matrixgame2.causal_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.matrixgame2.causal_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.matrixgame2.causal_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.matrixgame2.causal_model.ReplicatedLinear.weight_loader

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

fastvideo.models.dits.matrixgame2.causal_model.ReplicatedLinear.forward

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

fastvideo.models.dits.matrixgame2.causal_model.ReplicatedLinear.extra_repr

extra_repr() -> str

fastvideo.models.dits.matrixgame2.causal_model.ReplicatedLinear.get_shape_mapping

get_shape_mapping() -> dict

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

fastvideo.models.dits.matrixgame2.causal_model.ReplicatedLinear.reset_shape_tracking

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

fastvideo.models.dits.matrixgame2.causal_model.ReplicatedLinear.print_shape_summary

print_shape_summary() -> None

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

fastvideo.models.dits.matrixgame2.causal_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.matrixgame2.causal_model.MLP.forward

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

fastvideo.models.dits.matrixgame2.causal_model.get_rotary_pos_embed

get_rotary_pos_embed(rope_sizes, hidden_size, heads_num, rope_dim_list, rope_theta, theta_rescale_factor = 1.0, interpolation_factor = 1.0, shard_dim: int = 0, do_sp_sharding: bool = False, dtype: torch.dtype = torch.float32, start_frame: int = 0, use_real: bool = True) -> tuple[torch.Tensor, torch.Tensor]

Generate rotary positional embeddings for the given sizes.

Parameters

rope_sizes

Tuple of dimensions (t, h, w)

hidden_size

Hidden dimension size

heads_num

Number of attention heads

rope_dim_list

List of dimensions for each axis, or None

rope_theta

Base for frequency calculations

theta_rescale_factor

Rescale factor for theta. Defaults to 1.0

interpolation_factor

Factor to scale positions. Defaults to 1.0

shard_dim: int

Which dimension to shard for sequence parallelism. Defaults to 0.

do_sp_sharding: bool

Whether to shard the positional embeddings for sequence parallelism. Defaults to False.

use_real: bool

If True, output full head_dim for rotate_half style RoPE. Defaults to True.

Returns

tuple[torch.Tensor, torch.Tensor]

Tuple of (cos, sin) tensors for rotary embeddings. Shape [S, D] if use_real, [S, D/2] otherwise.

fastvideo.models.dits.matrixgame2.causal_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.matrixgame2.causal_model.PatchEmbed.forward

forward(x)

fastvideo.models.dits.matrixgame2.causal_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.matrixgame2.causal_model.TimestepEmbedder.forward

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

fastvideo.models.dits.matrixgame2.causal_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.matrixgame2.causal_model.ModulateProjection.forward

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

fastvideo.models.dits.matrixgame2.causal_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.matrixgame2.causal_model.relativistic_window_offsets

relativistic_window_offsets(local_end_index: int, num_new_tokens: int, max_attention_size: int) -> tuple[int, int, int]

Token offsets into a position-0 cos/sin table for one re-indexing step.

Returns (window_len, query_lo, query_hi): the cached window occupies table[0:window_len] and the query the tail table[query_lo:query_hi].

fastvideo.models.dits.matrixgame2.causal_model.BaseDiT

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

fastvideo.models.dits.matrixgame2.causal_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.matrixgame2.causal_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.matrixgame2.causal_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.matrixgame2.causal_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.matrixgame2.causal_model.WanT2VCrossAttention

WanT2VCrossAttention(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.matrixgame2.causal_model.WanT2VCrossAttention.forward

forward(x, context, context_lens, crossattn_cache = None)

Parameters

x: Tensor

Shape [B, L1, C]

context: Tensor

Shape [B, L2, C]

context_lens: Tensor

Shape [B]

fastvideo.models.dits.matrixgame2.causal_model.WanImageEmbedding

WanImageEmbedding(in_features: int, out_features: int)

fastvideo.models.dits.matrixgame2.causal_model.WanImageEmbedding.forward

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

fastvideo.models.dits.matrixgame2.causal_model.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.models.dits.matrixgame2.causal_model.ActionModule

ActionModule(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, local_attn_size = 6, sink_size = 0, blocks: list | None = None)

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

fastvideo.models.dits.matrixgame2.causal_model.ActionModule.patchify

patchify(x, patch_size)

x : (N C T H W)

fastvideo.models.dits.matrixgame2.causal_model.ActionModule.unpatchify

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

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

fastvideo.models.dits.matrixgame2.causal_model.ActionModule.get_rotary_pos_embed

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

fastvideo.models.dits.matrixgame2.causal_model.ActionModule.forward

forward(x: torch.Tensor, tt: int, th: int, tw: int, mouse_condition: torch.Tensor | None = None, keyboard_condition: torch.Tensor | None = None, block_mask_mouse: BlockMask | None = None, block_mask_keyboard: BlockMask | None = None, is_causal: bool = False, kv_cache_mouse: dict[str, torch.Tensor] | None = None, kv_cache_keyboard: dict[str, torch.Tensor] | None = None, start_frame: int = 0, use_rope_keyboard: bool = True, num_frame_per_block: int = 3)

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

fastvideo.models.dits.matrixgame2.causal_model.MatrixGame2CrossAttention

MatrixGame2CrossAttention(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.matrixgame2.causal_model.MatrixGame2CrossAttention.forward

forward(x, context, context_lens = None, crossattn_cache = None)

Parameters

x: Tensor

Shape [B, L1, C]

context: Tensor

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

context_lens: Tensor

Shape [B]

crossattn_cache: dict

Optional cache for k/v during inference

fastvideo.models.dits.matrixgame2.causal_model.retain_kv_with_sink

retain_kv_with_sink(k: torch.Tensor, v: torch.Tensor, target_len: int, sink_size: int) -> tuple[torch.Tensor, torch.Tensor]

fastvideo.models.dits.matrixgame2.causal_model.causal_rope_apply

causal_rope_apply(x: torch.Tensor, grid_sizes: tuple[int, int, int], freqs: torch.Tensor, start_frame: int = 0)
Source
def causal_rope_apply(x: torch.Tensor, grid_sizes: tuple[int, int, int], freqs: torch.Tensor, start_frame: int = 0):
    n, c = x.size(2), x.size(3) // 2

    # split freqs
    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)

    # loop over samples
    output = []
    f, h, w = grid_sizes

    for i in range(len(x)):
        seq_len = f * h * w

        # precompute multipliers
        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2))
        freqs_i = torch.cat(
            [
                freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1),
                freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
                freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
            ],
            dim=-1,
        ).reshape(seq_len, 1, -1)

        # apply rotary embedding
        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
        x_i = torch.cat([x_i, x[i, seq_len:]])

        # append to collection
        output.append(x_i)
    return torch.stack(output).type_as(x)

fastvideo.models.dits.matrixgame2.causal_model.rope_params

rope_params(max_seq_len, dim, theta = 10000)
Source
def rope_params(max_seq_len, dim, theta=10000):
    assert dim % 2 == 0
    freqs = torch.outer(
        torch.arange(max_seq_len),
        1.0 / torch.pow(theta,
                        torch.arange(0, dim, 2).to(torch.float64).div(dim)),
    )
    freqs = torch.polar(torch.ones_like(freqs), freqs)
    return freqs

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2TimeImageEmbedding

CausalMatrixGame2TimeImageEmbedding(dim: int, time_freq_dim: int, image_embed_dim: int | None = None)
Source
class CausalMatrixGame2TimeImageEmbedding(nn.Module):

    def __init__(
        self,
        dim: int,
        time_freq_dim: int,
        image_embed_dim: int | None = None,
    ):
        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")

        self.image_embedder = None
        if image_embed_dim is not None:
            self.image_embedder = WanImageEmbedding(image_embed_dim, dim)

    def forward(
        self,
        timestep: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        encoder_hidden_states_image: torch.Tensor | None = 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_image is not None:
            assert self.image_embedder is not None
            encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image)

        return temb, timestep_proj, None, encoder_hidden_states_image

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2TimeImageEmbedding.forward

forward(timestep: torch.Tensor, encoder_hidden_states: torch.Tensor, encoder_hidden_states_image: torch.Tensor | None = None, timestep_seq_len: int | None = None)
Source
def forward(
    self,
    timestep: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    encoder_hidden_states_image: torch.Tensor | None = 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_image is not None:
        assert self.image_embedder is not None
        encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image)

    return temb, timestep_proj, None, encoder_hidden_states_image

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2SelfAttention

CausalMatrixGame2SelfAttention(dim: int, num_heads: int, local_attn_size: int = -1, sink_size: int = 0, qk_norm: bool = True, eps: float = 1e-06, rope_cache_policy: str = 'absolute')
Source
class CausalMatrixGame2SelfAttention(nn.Module):

    def __init__(
        self,
        dim: int,
        num_heads: int,
        local_attn_size: int = -1,
        sink_size: int = 0,
        qk_norm: bool = True,
        eps: float = 1e-6,
        rope_cache_policy: str = "absolute",
    ) -> None:
        assert dim % num_heads == 0
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.local_attn_size = local_attn_size
        self.sink_size = sink_size
        self.qk_norm = qk_norm
        self.eps = eps
        self.rope_cache_policy = rope_cache_policy
        self._freqs_cache = None

        self.attn = LocalAttention(
            num_heads=num_heads,
            head_size=self.head_dim,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
            supported_attention_backends=(
                AttentionBackendEnum.FLASH_ATTN,
                AttentionBackendEnum.TORCH_SDPA,
            ),
        )

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        freqs_cis: tuple[torch.Tensor, torch.Tensor],
        block_mask: BlockMask,
        grid_sizes: tuple[int, int, int],
        kv_cache: dict | None = None,
        current_start: int = 0,
        cache_start: int | None = None,
    ):
        if cache_start is None:
            cache_start = current_start

        # Calculate start_frame for causal mode
        if kv_cache is not None:
            frame_seqlen = int(grid_sizes[1] * grid_sizes[2])
            start_frame = current_start // frame_seqlen
        else:
            start_frame = 0

        if self._freqs_cache is None or self._freqs_cache.device != q.device:
            d = self.head_dim
            self._freqs_cache = torch.cat(
                [
                    rope_params(1024, d - 4 * (d // 6)),
                    rope_params(1024, 2 * (d // 6)),
                    rope_params(1024, 2 * (d // 6)),
                ],
                dim=1,
            ).to(q.device)

        freqs = self._freqs_cache

        # relativistic defers roping until the cache window is known (and caches raw k)
        relativistic = self.rope_cache_policy == "relativistic" and kv_cache is not None
        if not relativistic:
            roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=start_frame).type_as(v)
            roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=start_frame).type_as(v)

        if kv_cache is None:
            padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
            padded_roped_query = torch.cat(
                [
                    roped_query,
                    torch.zeros(
                        [q.shape[0], padded_length, q.shape[2], q.shape[3]],
                        device=q.device,
                        dtype=v.dtype,
                    ),
                ],
                dim=1,
            )

            padded_roped_key = torch.cat(
                [
                    roped_key,
                    torch.zeros(
                        [k.shape[0], padded_length, k.shape[2], k.shape[3]],
                        device=k.device,
                        dtype=v.dtype,
                    ),
                ],
                dim=1,
            )

            padded_v = torch.cat(
                [
                    v,
                    torch.zeros(
                        [v.shape[0], padded_length, v.shape[2], v.shape[3]],
                        device=v.device,
                        dtype=v.dtype,
                    ),
                ],
                dim=1,
            )

            x = flex_attention(
                query=padded_roped_query.transpose(2, 1),
                key=padded_roped_key.transpose(2, 1),
                value=padded_v.transpose(2, 1),
                block_mask=block_mask,
            )[:, :, :-padded_length].transpose(2, 1)
        else:
            # Calculate frame_seqlen correctly from grid_sizes (single frame token count)
            if grid_sizes is not None:
                frame_seqlen = int(grid_sizes[1] * grid_sizes[2])
            else:
                # Fallback: assume q.shape[1] is single frame (shouldn't happen in causal mode)
                frame_seqlen = q.shape[1]
                logger.warning("grid_sizes not provided, using q.shape[1] as frame_seqlen")

            current_end = current_start + q.shape[1]
            sink_tokens = self.sink_size * frame_seqlen

            # Compute max_attention_size dynamically based on actual frame_seqlen
            max_attention_size = (15 * frame_seqlen if self.local_attn_size == -1 else self.local_attn_size *
                                  frame_seqlen)

            if torch.is_grad_enabled():
                kv_cache["k"] = kv_cache["k"].detach().clone()
                kv_cache["v"] = kv_cache["v"].detach().clone()
            else:
                kv_cache["k"] = kv_cache["k"].detach()
                kv_cache["v"] = kv_cache["v"].detach()

            kv_cache_size = kv_cache["k"].shape[1]
            num_new_tokens = q.shape[1]
            stored_key = k if relativistic else roped_key  # raw vs roped in cache
            global_end_index = (int(kv_cache["global_end_index"].item()) if isinstance(
                kv_cache["global_end_index"], torch.Tensor) else int(kv_cache["global_end_index"]))
            local_end_index_prev = (int(kv_cache["local_end_index"].item()) if isinstance(
                kv_cache["local_end_index"], torch.Tensor) else int(kv_cache["local_end_index"]))

            if (current_end > global_end_index) and (num_new_tokens + local_end_index_prev > kv_cache_size):
                num_evicted_tokens = (num_new_tokens + local_end_index_prev - kv_cache_size)
                num_rolled_tokens = (local_end_index_prev - num_evicted_tokens - sink_tokens)
                kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = kv_cache["k"][
                    :,
                    sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens,
                ].clone()
                kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = kv_cache["v"][
                    :,
                    sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens,
                ].clone()
                local_end_index = (local_end_index_prev + current_end - global_end_index - num_evicted_tokens)
                local_start_index = local_end_index - num_new_tokens
                kv_cache["k"][:, local_start_index:local_end_index] = stored_key
                kv_cache["v"][:, local_start_index:local_end_index] = v
            else:
                local_end_index = (local_end_index_prev + current_end - global_end_index)
                local_start_index = local_end_index - num_new_tokens
                kv_cache["k"][:, local_start_index:local_end_index] = stored_key
                kv_cache["v"][:, local_start_index:local_end_index] = v

            k_for_attn, v_for_attn = retain_kv_with_sink(
                kv_cache["k"][:, :local_end_index],
                kv_cache["v"][:, :local_end_index],
                min(local_end_index, max_attention_size),
                sink_tokens,
            )

            if relativistic:
                window_len, query_lo, _ = relativistic_window_offsets(local_end_index, num_new_tokens,
                                                                      max_attention_size)
                h, w = int(grid_sizes[1]), int(grid_sizes[2])
                window_grid = (window_len // frame_seqlen, h, w)
                roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=query_lo // frame_seqlen).type_as(v)
                k_for_attn = causal_rope_apply(k_for_attn, window_grid, freqs, start_frame=0).type_as(v)

            x = torch.nn.functional.scaled_dot_product_attention(
                roped_query.transpose(1, 2),
                k_for_attn.transpose(1, 2),
                v_for_attn.transpose(1, 2),
            ).transpose(1, 2)

            if isinstance(kv_cache["global_end_index"], torch.Tensor):
                kv_cache["global_end_index"].fill_(current_end)
            else:
                kv_cache["global_end_index"] = current_end
            if isinstance(kv_cache["local_end_index"], torch.Tensor):
                kv_cache["local_end_index"].fill_(local_end_index)
            else:
                kv_cache["local_end_index"] = local_end_index

        return x

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2SelfAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, grid_sizes: tuple[int, int, int], kv_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None)
Source
def forward(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    freqs_cis: tuple[torch.Tensor, torch.Tensor],
    block_mask: BlockMask,
    grid_sizes: tuple[int, int, int],
    kv_cache: dict | None = None,
    current_start: int = 0,
    cache_start: int | None = None,
):
    if cache_start is None:
        cache_start = current_start

    # Calculate start_frame for causal mode
    if kv_cache is not None:
        frame_seqlen = int(grid_sizes[1] * grid_sizes[2])
        start_frame = current_start // frame_seqlen
    else:
        start_frame = 0

    if self._freqs_cache is None or self._freqs_cache.device != q.device:
        d = self.head_dim
        self._freqs_cache = torch.cat(
            [
                rope_params(1024, d - 4 * (d // 6)),
                rope_params(1024, 2 * (d // 6)),
                rope_params(1024, 2 * (d // 6)),
            ],
            dim=1,
        ).to(q.device)

    freqs = self._freqs_cache

    # relativistic defers roping until the cache window is known (and caches raw k)
    relativistic = self.rope_cache_policy == "relativistic" and kv_cache is not None
    if not relativistic:
        roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=start_frame).type_as(v)
        roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=start_frame).type_as(v)

    if kv_cache is None:
        padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
        padded_roped_query = torch.cat(
            [
                roped_query,
                torch.zeros(
                    [q.shape[0], padded_length, q.shape[2], q.shape[3]],
                    device=q.device,
                    dtype=v.dtype,
                ),
            ],
            dim=1,
        )

        padded_roped_key = torch.cat(
            [
                roped_key,
                torch.zeros(
                    [k.shape[0], padded_length, k.shape[2], k.shape[3]],
                    device=k.device,
                    dtype=v.dtype,
                ),
            ],
            dim=1,
        )

        padded_v = torch.cat(
            [
                v,
                torch.zeros(
                    [v.shape[0], padded_length, v.shape[2], v.shape[3]],
                    device=v.device,
                    dtype=v.dtype,
                ),
            ],
            dim=1,
        )

        x = flex_attention(
            query=padded_roped_query.transpose(2, 1),
            key=padded_roped_key.transpose(2, 1),
            value=padded_v.transpose(2, 1),
            block_mask=block_mask,
        )[:, :, :-padded_length].transpose(2, 1)
    else:
        # Calculate frame_seqlen correctly from grid_sizes (single frame token count)
        if grid_sizes is not None:
            frame_seqlen = int(grid_sizes[1] * grid_sizes[2])
        else:
            # Fallback: assume q.shape[1] is single frame (shouldn't happen in causal mode)
            frame_seqlen = q.shape[1]
            logger.warning("grid_sizes not provided, using q.shape[1] as frame_seqlen")

        current_end = current_start + q.shape[1]
        sink_tokens = self.sink_size * frame_seqlen

        # Compute max_attention_size dynamically based on actual frame_seqlen
        max_attention_size = (15 * frame_seqlen if self.local_attn_size == -1 else self.local_attn_size *
                              frame_seqlen)

        if torch.is_grad_enabled():
            kv_cache["k"] = kv_cache["k"].detach().clone()
            kv_cache["v"] = kv_cache["v"].detach().clone()
        else:
            kv_cache["k"] = kv_cache["k"].detach()
            kv_cache["v"] = kv_cache["v"].detach()

        kv_cache_size = kv_cache["k"].shape[1]
        num_new_tokens = q.shape[1]
        stored_key = k if relativistic else roped_key  # raw vs roped in cache
        global_end_index = (int(kv_cache["global_end_index"].item()) if isinstance(
            kv_cache["global_end_index"], torch.Tensor) else int(kv_cache["global_end_index"]))
        local_end_index_prev = (int(kv_cache["local_end_index"].item()) if isinstance(
            kv_cache["local_end_index"], torch.Tensor) else int(kv_cache["local_end_index"]))

        if (current_end > global_end_index) and (num_new_tokens + local_end_index_prev > kv_cache_size):
            num_evicted_tokens = (num_new_tokens + local_end_index_prev - kv_cache_size)
            num_rolled_tokens = (local_end_index_prev - num_evicted_tokens - sink_tokens)
            kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = kv_cache["k"][
                :,
                sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens,
            ].clone()
            kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = kv_cache["v"][
                :,
                sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens,
            ].clone()
            local_end_index = (local_end_index_prev + current_end - global_end_index - num_evicted_tokens)
            local_start_index = local_end_index - num_new_tokens
            kv_cache["k"][:, local_start_index:local_end_index] = stored_key
            kv_cache["v"][:, local_start_index:local_end_index] = v
        else:
            local_end_index = (local_end_index_prev + current_end - global_end_index)
            local_start_index = local_end_index - num_new_tokens
            kv_cache["k"][:, local_start_index:local_end_index] = stored_key
            kv_cache["v"][:, local_start_index:local_end_index] = v

        k_for_attn, v_for_attn = retain_kv_with_sink(
            kv_cache["k"][:, :local_end_index],
            kv_cache["v"][:, :local_end_index],
            min(local_end_index, max_attention_size),
            sink_tokens,
        )

        if relativistic:
            window_len, query_lo, _ = relativistic_window_offsets(local_end_index, num_new_tokens,
                                                                  max_attention_size)
            h, w = int(grid_sizes[1]), int(grid_sizes[2])
            window_grid = (window_len // frame_seqlen, h, w)
            roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=query_lo // frame_seqlen).type_as(v)
            k_for_attn = causal_rope_apply(k_for_attn, window_grid, freqs, start_frame=0).type_as(v)

        x = torch.nn.functional.scaled_dot_product_attention(
            roped_query.transpose(1, 2),
            k_for_attn.transpose(1, 2),
            v_for_attn.transpose(1, 2),
        ).transpose(1, 2)

        if isinstance(kv_cache["global_end_index"], torch.Tensor):
            kv_cache["global_end_index"].fill_(current_end)
        else:
            kv_cache["global_end_index"] = current_end
        if isinstance(kv_cache["local_end_index"], torch.Tensor):
            kv_cache["local_end_index"].fill_(local_end_index)
        else:
            kv_cache["local_end_index"] = local_end_index

    return x

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2TransformerBlock

CausalMatrixGame2TransformerBlock(dim: int, ffn_dim: int, num_heads: int, local_attn_size: int = -1, sink_size: int = 0, qk_norm: str = 'rms_norm_across_heads', cross_attn_norm: bool = False, eps: float = 1e-06, added_kv_proj_dim: int | None = None, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '', action_config: dict | None = None, block_idx: int = 0, rope_cache_policy: str = 'absolute')
Source
class CausalMatrixGame2TransformerBlock(nn.Module):

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

        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 = CausalMatrixGame2SelfAttention(
            dim,
            num_heads,
            local_attn_size=local_attn_size,
            sink_size=sink_size,
            qk_norm=qk_norm,
            eps=eps,
            rope_cache_policy=rope_cache_policy,
        )
        self.hidden_dim = dim
        self.num_attention_heads = num_heads
        self.local_attn_size = local_attn_size
        self.sink_size = sink_size
        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":
            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,
        )

        if added_kv_proj_dim is not None:
            self.attn2 = MatrixGame2CrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
        else:
            self.attn2 = WanT2VCrossAttention(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,
        )

        if len(action_config) > 0 and block_idx in action_config.get("blocks", []):
            self.action_model = ActionModule(
                heads_num=action_config["heads_num"],
                hidden_size=action_config["hidden_size"],
                img_hidden_size=action_config["img_hidden_size"],
                mouse_hidden_dim=action_config.get("mouse_hidden_dim", 0),
                keyboard_hidden_dim=action_config["keyboard_hidden_dim"],
                mouse_dim_in=action_config.get("mouse_dim_in", 0),
                keyboard_dim_in=action_config["keyboard_dim_in"],
                enable_mouse=action_config["enable_mouse"],
                enable_keyboard=action_config["enable_keyboard"],
                windows_size=action_config["windows_size"],
                rope_theta=action_config["rope_theta"],
                rope_dim_list=action_config["rope_dim_list"],
                mouse_qk_dim_list=action_config.get("mouse_qk_dim_list", [8, 28, 28]),
                patch_size=action_config["patch_size"],
                local_attn_size=local_attn_size,
                sink_size=sink_size,
                qk_norm=action_config["qk_norm"],
                qkv_bias=action_config["qkv_bias"],
                vae_time_compression_ratio=action_config["vae_time_compression_ratio"],
            )
        else:
            self.action_model = None

        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)

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        temb: torch.Tensor,
        freqs_cis: tuple[torch.Tensor, torch.Tensor],
        block_mask: BlockMask,
        grid_sizes: tuple[int, int, int],
        mouse_cond: torch.Tensor | None = None,
        keyboard_cond: torch.Tensor | None = None,
        block_mask_mouse: BlockMask | None = None,
        block_mask_keyboard: BlockMask | None = None,
        num_frame_per_block: int = 1,
        use_rope_keyboard: bool = True,
        kv_cache: dict | None = None,
        kv_cache_mouse: dict | None = None,
        kv_cache_keyboard: dict | None = None,
        crossattn_cache: dict | None = None,
        current_start: int = 0,
        cache_start: int | None = None,
    ) -> torch.Tensor:
        if hidden_states.dim() == 4:
            hidden_states = hidden_states.squeeze(1)

        assert temb.ndim == 4
        num_frames = temb.shape[1]
        frame_seqlen = hidden_states.shape[1] // num_frames
        bs, seq_length, _ = hidden_states.shape
        orig_dtype = hidden_states.dtype
        # hidden_states = hidden_states.float()

        e = self.scale_shift_table + temb
        assert e.shape == (bs, num_frames, 6, self.hidden_dim)
        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (e.chunk(6, dim=2))

        norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
                              (1 + scale_msa) + shift_msa).flatten(1, 2)
        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))

        attn_output = self.attn1(
            query,
            key,
            value,
            freqs_cis,
            block_mask,
            grid_sizes,
            kv_cache,
            current_start,
            cache_start,
        )
        attn_output = attn_output.flatten(2)
        attn_output, _ = self.to_out(attn_output)
        attn_output = attn_output.squeeze(1)

        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
        norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
                                                                         null_shift, null_scale)

        # norm3 weights are loaded into self_attn_residual_norm.norm
        attn_output = self.attn2(
            norm_hidden_states,
            context=encoder_hidden_states,
            context_lens=None,
            crossattn_cache=crossattn_cache,
        )
        # residual connection
        hidden_states = hidden_states + attn_output

        if self.action_model is not None:
            if mouse_cond is not None or keyboard_cond is not None:
                start_frame = current_start // frame_seqlen

                hidden_states = self.action_model(
                    hidden_states,
                    grid_sizes[0],
                    grid_sizes[1],
                    grid_sizes[2],
                    mouse_cond,
                    keyboard_cond,
                    block_mask_mouse,
                    block_mask_keyboard,
                    is_causal=True,
                    kv_cache_mouse=kv_cache_mouse,
                    kv_cache_keyboard=kv_cache_keyboard,
                    start_frame=start_frame,
                    use_rope_keyboard=use_rope_keyboard,
                    num_frame_per_block=num_frame_per_block,
                )

        ff_output = self.ffn((self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
                              (1 + c_scale_msa) + c_shift_msa).flatten(1, 2))
        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)

        # Convert back to original dtype before return
        return hidden_states.to(orig_dtype)

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2TransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, grid_sizes: tuple[int, int, int], mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, block_mask_mouse: BlockMask | None = None, block_mask_keyboard: BlockMask | None = None, num_frame_per_block: int = 1, use_rope_keyboard: bool = True, kv_cache: dict | None = None, kv_cache_mouse: dict | None = None, kv_cache_keyboard: dict | None = None, crossattn_cache: dict | None = None, current_start: int = 0, cache_start: int | None = None) -> torch.Tensor
Source
def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    temb: torch.Tensor,
    freqs_cis: tuple[torch.Tensor, torch.Tensor],
    block_mask: BlockMask,
    grid_sizes: tuple[int, int, int],
    mouse_cond: torch.Tensor | None = None,
    keyboard_cond: torch.Tensor | None = None,
    block_mask_mouse: BlockMask | None = None,
    block_mask_keyboard: BlockMask | None = None,
    num_frame_per_block: int = 1,
    use_rope_keyboard: bool = True,
    kv_cache: dict | None = None,
    kv_cache_mouse: dict | None = None,
    kv_cache_keyboard: dict | None = None,
    crossattn_cache: dict | None = None,
    current_start: int = 0,
    cache_start: int | None = None,
) -> torch.Tensor:
    if hidden_states.dim() == 4:
        hidden_states = hidden_states.squeeze(1)

    assert temb.ndim == 4
    num_frames = temb.shape[1]
    frame_seqlen = hidden_states.shape[1] // num_frames
    bs, seq_length, _ = hidden_states.shape
    orig_dtype = hidden_states.dtype
    # hidden_states = hidden_states.float()

    e = self.scale_shift_table + temb
    assert e.shape == (bs, num_frames, 6, self.hidden_dim)
    shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (e.chunk(6, dim=2))

    norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
                          (1 + scale_msa) + shift_msa).flatten(1, 2)
    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))

    attn_output = self.attn1(
        query,
        key,
        value,
        freqs_cis,
        block_mask,
        grid_sizes,
        kv_cache,
        current_start,
        cache_start,
    )
    attn_output = attn_output.flatten(2)
    attn_output, _ = self.to_out(attn_output)
    attn_output = attn_output.squeeze(1)

    null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
    norm_hidden_states, hidden_states = self.self_attn_residual_norm(hidden_states, attn_output, gate_msa,
                                                                     null_shift, null_scale)

    # norm3 weights are loaded into self_attn_residual_norm.norm
    attn_output = self.attn2(
        norm_hidden_states,
        context=encoder_hidden_states,
        context_lens=None,
        crossattn_cache=crossattn_cache,
    )
    # residual connection
    hidden_states = hidden_states + attn_output

    if self.action_model is not None:
        if mouse_cond is not None or keyboard_cond is not None:
            start_frame = current_start // frame_seqlen

            hidden_states = self.action_model(
                hidden_states,
                grid_sizes[0],
                grid_sizes[1],
                grid_sizes[2],
                mouse_cond,
                keyboard_cond,
                block_mask_mouse,
                block_mask_keyboard,
                is_causal=True,
                kv_cache_mouse=kv_cache_mouse,
                kv_cache_keyboard=kv_cache_keyboard,
                start_frame=start_frame,
                use_rope_keyboard=use_rope_keyboard,
                num_frame_per_block=num_frame_per_block,
            )

    ff_output = self.ffn((self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
                          (1 + c_scale_msa) + c_shift_msa).flatten(1, 2))
    hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)

    # Convert back to original dtype before return
    return hidden_states.to(orig_dtype)

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2WanModel

CausalMatrixGame2WanModel(config: MatrixGame2WanVideoConfig, hf_config: dict[str, Any], **kwargs)
Source
class CausalMatrixGame2WanModel(BaseDiT):
    supports_action_input = True

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

    def __init__(
        self,
        config: MatrixGame2WanVideoConfig,
        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.attention_head_dim = config.attention_head_dim
        self.in_channels = config.in_channels
        self.out_channels = config.out_channels
        self.num_channels_latents = config.num_channels_latents
        self.patch_size = config.patch_size

        # Long tuning controls the causal window from YAML; consume it here
        # like Wan so train wrappers do not need runtime patching.
        arch_cfg = getattr(config, "arch_config", None)
        self.local_attn_size = (getattr(
            arch_cfg,
            "local_attn_size",
            getattr(config, "local_attn_size", -1),
        ) if arch_cfg else getattr(config, "local_attn_size", -1))
        self.sink_size = (getattr(arch_cfg, "sink_size", getattr(config, "sink_size", 0)) if arch_cfg else getattr(
            config, "sink_size", 0))
        if self.sink_size < 0:
            raise ValueError("sink_size must be non-negative")
        if self.local_attn_size != -1 and self.sink_size >= self.local_attn_size:
            raise ValueError("sink_size must be smaller than local_attn_size for "
                             "MatrixGame2 causal attention")
        self.rope_cache_policy = (getattr(arch_cfg, "rope_cache_policy", getattr(
            config, "rope_cache_policy", "absolute")) if arch_cfg else getattr(config, "rope_cache_policy", "absolute"))

        # 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 = CausalMatrixGame2TimeImageEmbedding(
            dim=inner_dim,
            time_freq_dim=config.freq_dim,
            image_embed_dim=config.image_dim,
        )

        # 2.1. Get action config
        arch_cfg = getattr(config, "arch_config", None)
        self.action_config = (getattr(arch_cfg, "action_config", {}) if arch_cfg is not None else {})

        # 3. Transformer blocks
        self.blocks = nn.ModuleList([
            CausalMatrixGame2TransformerBlock(
                inner_dim,
                config.ffn_dim,
                config.num_attention_heads,
                self.local_attn_size,
                self.sink_size,
                config.qk_norm,
                config.cross_attn_norm,
                config.eps,
                config.added_kv_proj_dim,
                self._supported_attention_backends,
                prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}",
                action_config=self.action_config,
                block_idx=i,
                rope_cache_policy=self.rope_cache_policy,
            ) 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,
        )
        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

        self.block_mask = None
        self.block_mask_keyboard = None
        self.block_mask_mouse = None
        self.use_rope_keyboard = True

        self.num_frame_per_block = (getattr(
            arch_cfg,
            "num_frames_per_block",
            getattr(config, "num_frame_per_block", 1),
        ) if arch_cfg else getattr(config, "num_frame_per_block", 1))

        self.__post_init__()

    def set_causal_attention_window(
        self,
        *,
        local_attn_size: int | None = None,
        sink_size: int | None = None,
    ) -> None:
        local_attn_size = (self.local_attn_size if local_attn_size is None else int(local_attn_size))
        sink_size = self.sink_size if sink_size is None else int(sink_size)
        if sink_size < 0:
            raise ValueError("sink_size must be non-negative")
        if local_attn_size != -1 and sink_size >= local_attn_size:
            raise ValueError("sink_size must be smaller than local_attn_size for "
                             "MatrixGame2 causal attention")

        self.local_attn_size = local_attn_size
        self.sink_size = sink_size
        self.block_mask = None
        self.block_mask_keyboard = None
        self.block_mask_mouse = None

        for block in self.blocks:
            if hasattr(block, "local_attn_size"):
                block.local_attn_size = local_attn_size
            if hasattr(block, "sink_size"):
                block.sink_size = sink_size

            attn1 = getattr(block, "attn1", None)
            if attn1 is not None:
                if hasattr(attn1, "local_attn_size"):
                    attn1.local_attn_size = local_attn_size
                if hasattr(attn1, "sink_size"):
                    attn1.sink_size = sink_size

            action_model = getattr(block, "action_model", None)
            if action_model is not None:
                if hasattr(action_model, "local_attn_size"):
                    action_model.local_attn_size = local_attn_size
                if hasattr(action_model, "sink_size"):
                    action_model.sink_size = sink_size

    def set_sink_size(self, sink_size: int) -> None:
        self.set_causal_attention_window(sink_size=sink_size)

    @staticmethod
    def _prepare_blockwise_causal_attn_mask(
        device: torch.device | str,
        num_frames: int = 9,
        frame_seqlen: int = 880,
        num_frame_per_block: int = 1,
        local_attn_size: int = -1,
        sink_size: int = 0,
    ) -> BlockMask:
        total_length = num_frames * frame_seqlen
        padded_length = math.ceil(total_length / 128) * 128 - total_length

        ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)

        frame_indices = torch.arange(
            start=0,
            end=total_length,
            step=frame_seqlen * num_frame_per_block,
            device=device,
        )

        for tmp in frame_indices:
            ends[tmp:tmp + frame_seqlen * num_frame_per_block] = (tmp + frame_seqlen * num_frame_per_block)

        def attention_mask(b, h, q_idx, kv_idx):
            if local_attn_size == -1:
                return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)
            else:
                return ((kv_idx < ends[q_idx])
                        & ((kv_idx < sink_size * frame_seqlen)
                           | (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen)))) | (q_idx == kv_idx)

        block_mask = create_block_mask(
            attention_mask,
            B=None,
            H=None,
            Q_LEN=total_length + padded_length,
            KV_LEN=total_length + padded_length,
            _compile=False,
            device=device,
        )

        if not dist.is_initialized() or dist.get_rank() == 0:
            logger.info(
                "cache a block wise causal mask with block size of %s frames",
                num_frame_per_block,
            )

        return block_mask

    @staticmethod
    def _prepare_blockwise_causal_attn_mask_keyboard(
        device: torch.device | str,
        num_frames: int = 9,
        frame_seqlen: int = 880,
        num_frame_per_block: int = 1,
        local_attn_size: int = -1,
        sink_size: int = 0,
    ) -> BlockMask:
        total_length2 = num_frames * frame_seqlen
        padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
        padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames
        ends2 = torch.zeros(total_length2 + padded_length2, device=device, dtype=torch.long)

        frame_indices2 = torch.arange(
            start=0,
            end=total_length2,
            step=frame_seqlen * num_frame_per_block,
            device=device,
        )
        cnt = num_frame_per_block
        for tmp in frame_indices2:
            ends2[tmp:tmp + frame_seqlen * num_frame_per_block] = cnt
            cnt += num_frame_per_block

        def attention_mask2(b, h, q_idx, kv_idx):
            if local_attn_size == -1:
                return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx)
            else:
                return ((kv_idx < ends2[q_idx])
                        & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))) | (q_idx == kv_idx)

        block_mask2 = create_block_mask(
            attention_mask2,
            B=None,
            H=None,
            Q_LEN=total_length2 + padded_length2,
            KV_LEN=num_frames + padded_length_kv2,
            _compile=False,
            device=device,
        )

        if not dist.is_initialized() or dist.get_rank() == 0:
            logger.info(
                "cache a block wise causal mask for keyboard with block size of %s frames",
                num_frame_per_block,
            )

        return block_mask2

    @staticmethod
    def _prepare_blockwise_causal_attn_mask_action(
        device: torch.device | str,
        num_frames: int = 9,
        frame_seqlen: int = 1,
        num_frame_per_block: int = 1,
        local_attn_size: int = -1,
        sink_size: int = 0,
    ) -> BlockMask:
        total_length2 = num_frames * frame_seqlen
        padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2
        padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames
        ends2 = torch.zeros(total_length2 + padded_length2, device=device, dtype=torch.long)

        frame_indices2 = torch.arange(
            start=0,
            end=total_length2,
            step=frame_seqlen * num_frame_per_block,
            device=device,
        )
        cnt = num_frame_per_block
        for tmp in frame_indices2:
            ends2[tmp:tmp + frame_seqlen * num_frame_per_block] = cnt
            cnt += num_frame_per_block

        def attention_mask2(b, h, q_idx, kv_idx):
            if local_attn_size == -1:
                return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx)
            else:
                return ((kv_idx < ends2[q_idx])
                        & ((kv_idx < sink_size) | (kv_idx >= (ends2[q_idx] - local_attn_size)))) | (q_idx == kv_idx)

        block_mask2 = create_block_mask(
            attention_mask2,
            B=None,
            H=None,
            Q_LEN=total_length2 + padded_length2,
            KV_LEN=num_frames + padded_length_kv2,
            _compile=False,
            device=device,
        )

        if not dist.is_initialized() or dist.get_rank() == 0:
            logger.info(
                "cache a block wise causal mask for action with block size of %s frames",
                num_frame_per_block,
            )

        return block_mask2

    def _forward_inference(
        self,
        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,
        mouse_cond: torch.Tensor | None = None,
        keyboard_cond: torch.Tensor | None = None,
        kv_cache: dict = None,
        kv_cache_mouse: dict = None,
        kv_cache_keyboard: dict = None,
        crossattn_cache: dict = None,
        current_start: int = 0,
        cache_start: int = 0,
        start_frame: int = 0,
        **kwargs,
    ) -> torch.Tensor:
        # Extract num_frame_per_block from kwargs
        effective_num_frame_per_block = kwargs.pop("num_frame_per_block", self.num_frame_per_block)

        if mouse_cond is not None or keyboard_cond is not None:
            assert (self.action_config is not None and len(self.action_config) > 0)

        if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
        if isinstance(encoder_hidden_states_image, list):
            encoder_hidden_states_image = (encoder_hidden_states_image[0]
                                           if len(encoder_hidden_states_image) > 0 else None)

        ctx = get_forward_context()
        batch = getattr(ctx, "forward_batch", None)

        if (batch is not None and getattr(batch, "image_latent", None) is not None):
            image_latent = batch.image_latent
            if (isinstance(image_latent, torch.Tensor) and image_latent.ndim == 5 and hidden_states.ndim == 5):
                _, _, num_frames, _, _ = hidden_states.shape
                start = start_frame
                end = start + num_frames
                cond = image_latent
                if cond.shape[2] >= end:
                    cond = cond[:, :, start:end]
                elif cond.shape[2] > start:
                    cond = cond[:, :, start:]
                    pad_frames = num_frames - cond.shape[2]
                    if pad_frames > 0:
                        pad = torch.zeros(
                            cond.shape[0],
                            cond.shape[1],
                            pad_frames,
                            cond.shape[3],
                            cond.shape[4],
                            device=cond.device,
                            dtype=cond.dtype,
                        )
                        cond = torch.cat([cond, pad], dim=2)
                else:
                    pad = torch.zeros(
                        cond.shape[0],
                        cond.shape[1],
                        num_frames,
                        cond.shape[3],
                        cond.shape[4],
                        device=cond.device,
                        dtype=cond.dtype,
                    )
                    cond = pad
                hidden_states = torch.cat([hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1)

        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

        d = self.hidden_size // self.num_attention_heads
        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
        freqs_cos, freqs_sin = get_rotary_pos_embed(
            (
                post_patch_num_frames * get_sp_world_size(),
                post_patch_height,
                post_patch_width,
            ),
            self.hidden_size,
            self.num_attention_heads,
            rope_dim_list,
            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
            rope_theta=10000,
            start_frame=start_frame,
        )
        freqs_cos = freqs_cos.to(hidden_states.device)
        freqs_sin = freqs_sin.to(hidden_states.device)
        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None

        hidden_states = self.patch_embedding(hidden_states)
        grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)
        hidden_states = hidden_states.flatten(2).transpose(1, 2)

        if timestep.dim() == 2:
            timestep = timestep.flatten()

        (
            temb,
            timestep_proj,
            encoder_hidden_states_text,
            encoder_hidden_states_image,
        ) = self.condition_embedder(timestep, encoder_hidden_states, encoder_hidden_states_image)
        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0,
                                                                                    sizes=(batch_size,
                                                                                           post_patch_num_frames))

        if encoder_hidden_states_text is not None:
            if isinstance(encoder_hidden_states_text, list):
                encoder_hidden_states_text = encoder_hidden_states_text[0]
            elif encoder_hidden_states_text.ndim == 2:
                encoder_hidden_states_text = (encoder_hidden_states_text.unsqueeze(0))

        if encoder_hidden_states_image is not None:
            if encoder_hidden_states_text is not None:
                encoder_hidden_states = torch.concat(
                    [encoder_hidden_states_image, encoder_hidden_states_text],
                    dim=1,
                )
            else:
                encoder_hidden_states = encoder_hidden_states_image

        block_mask = self._prepare_blockwise_causal_attn_mask(
            device=hidden_states.device,
            num_frames=num_frames,
            frame_seqlen=post_patch_height * post_patch_width,
            num_frame_per_block=self.num_frame_per_block,
            local_attn_size=self.local_attn_size,
            sink_size=self.sink_size,
        )
        if self.use_rope_keyboard:
            block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action(
                device=hidden_states.device,
                num_frames=num_frames,
                frame_seqlen=1,
                num_frame_per_block=self.num_frame_per_block,
                local_attn_size=self.local_attn_size,
                sink_size=self.sink_size,
            )
        else:
            block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
                device=hidden_states.device,
                num_frames=num_frames,
                frame_seqlen=post_patch_height * post_patch_width,
                num_frame_per_block=self.num_frame_per_block,
                local_attn_size=self.local_attn_size,
                sink_size=self.sink_size,
            )
        block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
            device=hidden_states.device,
            num_frames=num_frames,
            frame_seqlen=1,
            num_frame_per_block=self.num_frame_per_block,
            local_attn_size=self.local_attn_size,
            sink_size=self.sink_size,
        )
        if kv_cache is None:
            kv_cache = [None] * len(self.blocks)
        if kv_cache_mouse is None:
            kv_cache_mouse = [None] * len(self.blocks)
        if kv_cache_keyboard is None:
            kv_cache_keyboard = [None] * len(self.blocks)
        if crossattn_cache is None:
            crossattn_cache = [None] * len(self.blocks)

        for block_index, block in enumerate(self.blocks):
            causal_kwargs = {
                "kv_cache": kv_cache[block_index],
                "kv_cache_mouse": kv_cache_mouse[block_index] if kv_cache_mouse else None,
                "kv_cache_keyboard": kv_cache_keyboard[block_index] if kv_cache_keyboard else None,
                "crossattn_cache": crossattn_cache[block_index] if crossattn_cache else None,
                "current_start": current_start,
                "cache_start": cache_start,
            }
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                causal_kwargs["num_frame_per_block"] = effective_num_frame_per_block
                hidden_states = self._gradient_checkpointing_func(block, hidden_states, **causal_kwargs)
            else:
                hidden_states = block(
                    hidden_states,
                    encoder_hidden_states,
                    timestep_proj,
                    freqs_cis,
                    block_mask=block_mask,
                    grid_sizes=grid_sizes,
                    mouse_cond=mouse_cond,
                    keyboard_cond=keyboard_cond,
                    block_mask_mouse=block_mask_mouse,
                    block_mask_keyboard=block_mask_keyboard,
                    num_frame_per_block=effective_num_frame_per_block,
                    use_rope_keyboard=self.use_rope_keyboard,
                    **causal_kwargs,
                )

        temb = temb.unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)).unsqueeze(2)
        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2)
        hidden_states = self.norm_out(hidden_states, shift, scale)
        hidden_states = self.proj_out(hidden_states)

        hidden_states = self.unpatchify(hidden_states, grid_sizes)

        return hidden_states

    def _forward_train(
        self,
        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,
        mouse_cond: torch.Tensor | None = None,
        keyboard_cond: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        if mouse_cond is not None or keyboard_cond is not None:
            assert (self.action_config is not None and len(self.action_config) > 0)

        if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
        if isinstance(encoder_hidden_states_image, list):
            encoder_hidden_states_image = (encoder_hidden_states_image[0]
                                           if len(encoder_hidden_states_image) > 0 else None)

        ctx = get_forward_context()
        batch = getattr(ctx, "forward_batch", None)
        if (batch is not None and getattr(batch, "image_latent", None) is not None):
            image_latent = batch.image_latent
            if (isinstance(image_latent, torch.Tensor) and image_latent.ndim == 5 and hidden_states.ndim == 5):
                _, _, num_frames, _, _ = hidden_states.shape
                cond = image_latent
                if cond.shape[2] >= num_frames:
                    cond = cond[:, :, :num_frames]
                else:
                    pad_frames = num_frames - cond.shape[2]
                    pad = torch.zeros(
                        cond.shape[0],
                        cond.shape[1],
                        pad_frames,
                        cond.shape[3],
                        cond.shape[4],
                        device=cond.device,
                        dtype=cond.dtype,
                    )
                    cond = torch.cat([cond, pad], dim=2)
                hidden_states = torch.cat([hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1)

        batch_size, num_channels, num_frames, height, width = (hidden_states.shape)
        p_t, p_h, p_w = self.patch_size
        post_patch_num_frames: int = num_frames // p_t
        post_patch_height: int = height // p_h
        post_patch_width: int = width // p_w

        d = self.hidden_size // self.num_attention_heads
        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
        freqs_cos, freqs_sin = get_rotary_pos_embed(
            (
                post_patch_num_frames * get_sp_world_size(),
                post_patch_height,
                post_patch_width,
            ),
            self.hidden_size,
            self.num_attention_heads,
            rope_dim_list,
            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
            rope_theta=10000,
        )
        freqs_cos = freqs_cos.to(hidden_states.device)
        freqs_sin = freqs_sin.to(hidden_states.device)
        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None

        if self.block_mask is None:
            self.block_mask = self._prepare_blockwise_causal_attn_mask(
                device=hidden_states.device,
                num_frames=num_frames,
                frame_seqlen=post_patch_height * post_patch_width,
                num_frame_per_block=self.num_frame_per_block,
                local_attn_size=self.local_attn_size,
                sink_size=self.sink_size,
            )
        if self.block_mask_keyboard is None:
            if self.use_rope_keyboard:
                self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action(
                    device=hidden_states.device,
                    num_frames=num_frames,
                    frame_seqlen=1,
                    num_frame_per_block=self.num_frame_per_block,
                    local_attn_size=self.local_attn_size,
                    sink_size=self.sink_size,
                )
            else:
                self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard(
                    device=hidden_states.device,
                    num_frames=num_frames,
                    frame_seqlen=post_patch_height * post_patch_width,
                    num_frame_per_block=self.num_frame_per_block,
                    local_attn_size=self.local_attn_size,
                    sink_size=self.sink_size,
                )
        if self.block_mask_mouse is None:
            self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action(
                device=hidden_states.device,
                num_frames=num_frames,
                frame_seqlen=1,
                num_frame_per_block=self.num_frame_per_block,
                local_attn_size=self.local_attn_size,
                sink_size=self.sink_size,
            )

        hidden_states = self.patch_embedding(hidden_states)
        grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)
        hidden_states = hidden_states.flatten(2).transpose(1, 2)

        if timestep.dim() == 2:
            timestep = timestep.flatten()

        (
            temb,
            timestep_proj,
            encoder_hidden_states_text,
            encoder_hidden_states_image,
        ) = self.condition_embedder(timestep, encoder_hidden_states, encoder_hidden_states_image)
        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0,
                                                                                    sizes=(batch_size,
                                                                                           post_patch_num_frames))

        if encoder_hidden_states_text is not None:
            if isinstance(encoder_hidden_states_text, list):
                encoder_hidden_states_text = encoder_hidden_states_text[0]
            elif encoder_hidden_states_text.ndim == 2:
                encoder_hidden_states_text = (encoder_hidden_states_text.unsqueeze(0))

        if encoder_hidden_states_image is not None:
            if encoder_hidden_states_text is not None:
                encoder_hidden_states = torch.concat(
                    [encoder_hidden_states_image, encoder_hidden_states_text],
                    dim=1,
                )
            else:
                encoder_hidden_states = encoder_hidden_states_image
        else:
            encoder_hidden_states = encoder_hidden_states_text

        kwargs = dict(
            e=timestep_proj,
            grid_sizes=grid_sizes,
            freqs_cis=freqs_cis,
            encoder_hidden_states=encoder_hidden_states,
            mouse_cond=mouse_cond,
            keyboard_cond=keyboard_cond,
            block_mask=self.block_mask,
            block_mask_mouse=self.block_mask_mouse,
            block_mask_keyboard=self.block_mask_keyboard,
            use_rope_keyboard=self.use_rope_keyboard,
            num_frame_per_block=self.num_frame_per_block,
        )

        # Initialize kv_cache, kv_cache_mouse, kv_cache_keyboard, crossattn_cache if not provided
        kv_cache = kwargs.pop("kv_cache", [None] * len(self.blocks))
        kv_cache_mouse = kwargs.pop("kv_cache_mouse", [None] * len(self.blocks))
        kv_cache_keyboard = kwargs.pop("kv_cache_keyboard", [None] * len(self.blocks))
        crossattn_cache = kwargs.pop("crossattn_cache", [None] * len(self.blocks))

        for i, block in enumerate(self.blocks):
            block_crossattn_cache = None
            if crossattn_cache is not None:
                if crossattn_cache[i] is None:
                    crossattn_cache[i] = {"is_init": False}
                block_crossattn_cache = crossattn_cache[i]

            if torch.is_grad_enabled() and self.gradient_checkpointing:
                hidden_states = self._gradient_checkpointing_func(
                    block,
                    hidden_states,
                    encoder_hidden_states,
                    timestep_proj,
                    freqs_cis,
                    kv_cache=kv_cache[i],
                    kv_cache_mouse=kv_cache_mouse[i],
                    kv_cache_keyboard=kv_cache_keyboard[i],
                    crossattn_cache=block_crossattn_cache,
                    block_mask=self.block_mask,
                    grid_sizes=grid_sizes,
                    mouse_cond=mouse_cond,
                    keyboard_cond=keyboard_cond,
                    block_mask_mouse=self.block_mask_mouse,
                    block_mask_keyboard=self.block_mask_keyboard,
                    num_frame_per_block=self.num_frame_per_block,
                    use_rope_keyboard=self.use_rope_keyboard,
                )
            else:
                hidden_states = block(
                    hidden_states,
                    encoder_hidden_states,
                    timestep_proj,
                    freqs_cis,
                    kv_cache=kv_cache[i],
                    kv_cache_mouse=kv_cache_mouse[i],
                    kv_cache_keyboard=kv_cache_keyboard[i],
                    crossattn_cache=block_crossattn_cache,
                    block_mask=self.block_mask,
                    grid_sizes=grid_sizes,
                    mouse_cond=mouse_cond,
                    keyboard_cond=keyboard_cond,
                    block_mask_mouse=self.block_mask_mouse,
                    block_mask_keyboard=self.block_mask_keyboard,
                    num_frame_per_block=self.num_frame_per_block,
                    use_rope_keyboard=self.use_rope_keyboard,
                )

        temb = temb.unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)).unsqueeze(2)
        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2)
        hidden_states = self.norm_out(hidden_states, shift, scale)
        hidden_states = self.proj_out(hidden_states)

        hidden_states = self.unpatchify(hidden_states, grid_sizes)

        return hidden_states

    def forward(self, *args, **kwargs):
        if kwargs.get("kv_cache") is not None:
            return self._forward_inference(*args, **kwargs)
        else:
            return self._forward_train(*args, **kwargs)

    def unpatchify(self, x: torch.Tensor, grid_sizes: tuple[int, int, int]) -> torch.Tensor:
        c = self.proj_out.out_features // math.prod(self.patch_size)
        p_t, p_h, p_w = self.patch_size
        f, h, w = grid_sizes

        x = x[:, :f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c)
        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6)
        x = x.reshape(-1, c, f * p_t, h * p_h, w * p_w)
        return x

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

        nn.init.xavier_uniform_(self.patch_embedding.proj.weight.flatten(1))
        for m in self.condition_embedder.time_embedder.mlp.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=0.02)

        nn.init.zeros_(self.proj_out.weight)

        for block in self.blocks:
            if block.action_model is not None:
                try:
                    nn.init.zeros_(block.action_model.proj_mouse.weight)
                    if block.action_model.proj_mouse.bias is not None:
                        nn.init.zeros_(block.action_model.proj_mouse.bias)
                    nn.init.zeros_(block.action_model.proj_keyboard.weight)
                    if block.action_model.proj_keyboard.bias is not None:
                        nn.init.zeros_(block.action_model.proj_keyboard.bias)
                except AttributeError:
                    pass

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2WanModel.set_causal_attention_window

set_causal_attention_window(*, local_attn_size: int | None = None, sink_size: int | None = None) -> None
Source
def set_causal_attention_window(
    self,
    *,
    local_attn_size: int | None = None,
    sink_size: int | None = None,
) -> None:
    local_attn_size = (self.local_attn_size if local_attn_size is None else int(local_attn_size))
    sink_size = self.sink_size if sink_size is None else int(sink_size)
    if sink_size < 0:
        raise ValueError("sink_size must be non-negative")
    if local_attn_size != -1 and sink_size >= local_attn_size:
        raise ValueError("sink_size must be smaller than local_attn_size for "
                         "MatrixGame2 causal attention")

    self.local_attn_size = local_attn_size
    self.sink_size = sink_size
    self.block_mask = None
    self.block_mask_keyboard = None
    self.block_mask_mouse = None

    for block in self.blocks:
        if hasattr(block, "local_attn_size"):
            block.local_attn_size = local_attn_size
        if hasattr(block, "sink_size"):
            block.sink_size = sink_size

        attn1 = getattr(block, "attn1", None)
        if attn1 is not None:
            if hasattr(attn1, "local_attn_size"):
                attn1.local_attn_size = local_attn_size
            if hasattr(attn1, "sink_size"):
                attn1.sink_size = sink_size

        action_model = getattr(block, "action_model", None)
        if action_model is not None:
            if hasattr(action_model, "local_attn_size"):
                action_model.local_attn_size = local_attn_size
            if hasattr(action_model, "sink_size"):
                action_model.sink_size = sink_size

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2WanModel.set_sink_size

set_sink_size(sink_size: int) -> None
Source
def set_sink_size(self, sink_size: int) -> None:
    self.set_causal_attention_window(sink_size=sink_size)

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2WanModel.forward

forward(*args, **kwargs)
Source
def forward(self, *args, **kwargs):
    if kwargs.get("kv_cache") is not None:
        return self._forward_inference(*args, **kwargs)
    else:
        return self._forward_train(*args, **kwargs)

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2WanModel.unpatchify

unpatchify(x: torch.Tensor, grid_sizes: tuple[int, int, int]) -> torch.Tensor
Source
def unpatchify(self, x: torch.Tensor, grid_sizes: tuple[int, int, int]) -> torch.Tensor:
    c = self.proj_out.out_features // math.prod(self.patch_size)
    p_t, p_h, p_w = self.patch_size
    f, h, w = grid_sizes

    x = x[:, :f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c)
    x = x.permute(0, 7, 1, 4, 2, 5, 3, 6)
    x = x.reshape(-1, c, f * p_t, h * p_h, w * p_w)
    return x

fastvideo.models.dits.matrixgame2.causal_model.CausalMatrixGame2WanModel.init_weights

init_weights()
Source
def init_weights(self):
    for m in self.modules():
        if isinstance(m, nn.Linear):
            nn.init.xavier_uniform_(m.weight)
            if m.bias is not None:
                nn.init.zeros_(m.bias)

    nn.init.xavier_uniform_(self.patch_embedding.proj.weight.flatten(1))
    for m in self.condition_embedder.time_embedder.mlp.modules():
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)

    nn.init.zeros_(self.proj_out.weight)

    for block in self.blocks:
        if block.action_model is not None:
            try:
                nn.init.zeros_(block.action_model.proj_mouse.weight)
                if block.action_model.proj_mouse.bias is not None:
                    nn.init.zeros_(block.action_model.proj_mouse.bias)
                nn.init.zeros_(block.action_model.proj_keyboard.weight)
                if block.action_model.proj_keyboard.bias is not None:
                    nn.init.zeros_(block.action_model.proj_keyboard.bias)
            except AttributeError:
                pass