causal_model

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)</code></pre></details>
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=&quot;silu&quot;)
    self.time_modulation = ModulateProjection(dim, factor=6, act_layer=&quot;silu&quot;)

    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</code></pre></details>
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</code></pre></details>
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 = &quot;absolute&quot;,
) -&gt; 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&#91;torch.Tensor, torch.Tensor&#93;,
    block_mask: BlockMask,
    grid_sizes: tuple&#91;int, int, int&#93;,
    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&#91;1&#93; * grid_sizes&#91;2&#93;)
        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(
            &#91;
                rope_params(1024, d - 4 * (d // 6)),
                rope_params(1024, 2 * (d // 6)),
                rope_params(1024, 2 * (d // 6)),
            &#93;,
            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 == &quot;relativistic&quot; 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&#91;1&#93; / 128) * 128 - q.shape&#91;1&#93;
        padded_roped_query = torch.cat(
            &#91;
                roped_query,
                torch.zeros(
                    &#91;q.shape&#91;0&#93;, padded_length, q.shape&#91;2&#93;, q.shape&#91;3&#93;&#93;,
                    device=q.device,
                    dtype=v.dtype,
                ),
            &#93;,
            dim=1,
        )

        padded_roped_key = torch.cat(
            &#91;
                roped_key,
                torch.zeros(
                    &#91;k.shape&#91;0&#93;, padded_length, k.shape&#91;2&#93;, k.shape&#91;3&#93;&#93;,
                    device=k.device,
                    dtype=v.dtype,
                ),
            &#93;,
            dim=1,
        )

        padded_v = torch.cat(
            &#91;
                v,
                torch.zeros(
                    &#91;v.shape&#91;0&#93;, padded_length, v.shape&#91;2&#93;, v.shape&#91;3&#93;&#93;,
                    device=v.device,
                    dtype=v.dtype,
                ),
            &#93;,
            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,
        )&#91;:, :, :-padded_length&#93;.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&#91;1&#93; * grid_sizes&#91;2&#93;)
        else:
            # Fallback: assume q.shape&#91;1&#93; is single frame (shouldn&#x27;t happen in causal mode)
            frame_seqlen = q.shape&#91;1&#93;
            logger.warning(&quot;grid_sizes not provided, using q.shape&#91;1&#93; as frame_seqlen&quot;)

        current_end = current_start + q.shape&#91;1&#93;
        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&#91;&quot;k&quot;&#93; = kv_cache&#91;&quot;k&quot;&#93;.detach().clone()
            kv_cache&#91;&quot;v&quot;&#93; = kv_cache&#91;&quot;v&quot;&#93;.detach().clone()
        else:
            kv_cache&#91;&quot;k&quot;&#93; = kv_cache&#91;&quot;k&quot;&#93;.detach()
            kv_cache&#91;&quot;v&quot;&#93; = kv_cache&#91;&quot;v&quot;&#93;.detach()

        kv_cache_size = kv_cache&#91;&quot;k&quot;&#93;.shape&#91;1&#93;
        num_new_tokens = q.shape&#91;1&#93;
        stored_key = k if relativistic else roped_key  # raw vs roped in cache
        global_end_index = (int(kv_cache&#91;&quot;global_end_index&quot;&#93;.item()) if isinstance(
            kv_cache&#91;&quot;global_end_index&quot;&#93;, torch.Tensor) else int(kv_cache&#91;&quot;global_end_index&quot;&#93;))
        local_end_index_prev = (int(kv_cache&#91;&quot;local_end_index&quot;&#93;.item()) if isinstance(
            kv_cache&#91;&quot;local_end_index&quot;&#93;, torch.Tensor) else int(kv_cache&#91;&quot;local_end_index&quot;&#93;))

        if (current_end &gt; global_end_index) and (num_new_tokens + local_end_index_prev &gt; 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&#91;&quot;k&quot;&#93;&#91;:, sink_tokens:sink_tokens + num_rolled_tokens&#93; = kv_cache&#91;&quot;k&quot;&#93;&#91;
                :,
                sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens,
            &#93;.clone()
            kv_cache&#91;&quot;v&quot;&#93;&#91;:, sink_tokens:sink_tokens + num_rolled_tokens&#93; = kv_cache&#91;&quot;v&quot;&#93;&#91;
                :,
                sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens,
            &#93;.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&#91;&quot;k&quot;&#93;&#91;:, local_start_index:local_end_index&#93; = stored_key
            kv_cache&#91;&quot;v&quot;&#93;&#91;:, local_start_index:local_end_index&#93; = 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&#91;&quot;k&quot;&#93;&#91;:, local_start_index:local_end_index&#93; = stored_key
            kv_cache&#91;&quot;v&quot;&#93;&#91;:, local_start_index:local_end_index&#93; = v

        k_for_attn, v_for_attn = retain_kv_with_sink(
            kv_cache&#91;&quot;k&quot;&#93;&#91;:, :local_end_index&#93;,
            kv_cache&#91;&quot;v&quot;&#93;&#91;:, :local_end_index&#93;,
            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&#91;1&#93;), int(grid_sizes&#91;2&#93;)
            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&#91;&quot;global_end_index&quot;&#93;, torch.Tensor):
            kv_cache&#91;&quot;global_end_index&quot;&#93;.fill_(current_end)
        else:
            kv_cache&#91;&quot;global_end_index&quot;&#93; = current_end
        if isinstance(kv_cache&#91;&quot;local_end_index&quot;&#93;, torch.Tensor):
            kv_cache&#91;&quot;local_end_index&quot;&#93;.fill_(local_end_index)
        else:
            kv_cache&#91;&quot;local_end_index&quot;&#93; = local_end_index

    return x</code></pre></details>
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&#91;1&#93; * grid_sizes&#91;2&#93;)
    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(
        &#91;
            rope_params(1024, d - 4 * (d // 6)),
            rope_params(1024, 2 * (d // 6)),
            rope_params(1024, 2 * (d // 6)),
        &#93;,
        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 == &quot;relativistic&quot; 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&#91;1&#93; / 128) * 128 - q.shape&#91;1&#93;
    padded_roped_query = torch.cat(
        &#91;
            roped_query,
            torch.zeros(
                &#91;q.shape&#91;0&#93;, padded_length, q.shape&#91;2&#93;, q.shape&#91;3&#93;&#93;,
                device=q.device,
                dtype=v.dtype,
            ),
        &#93;,
        dim=1,
    )

    padded_roped_key = torch.cat(
        &#91;
            roped_key,
            torch.zeros(
                &#91;k.shape&#91;0&#93;, padded_length, k.shape&#91;2&#93;, k.shape&#91;3&#93;&#93;,
                device=k.device,
                dtype=v.dtype,
            ),
        &#93;,
        dim=1,
    )

    padded_v = torch.cat(
        &#91;
            v,
            torch.zeros(
                &#91;v.shape&#91;0&#93;, padded_length, v.shape&#91;2&#93;, v.shape&#91;3&#93;&#93;,
                device=v.device,
                dtype=v.dtype,
            ),
        &#93;,
        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,
    )&#91;:, :, :-padded_length&#93;.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&#91;1&#93; * grid_sizes&#91;2&#93;)
    else:
        # Fallback: assume q.shape&#91;1&#93; is single frame (shouldn&#x27;t happen in causal mode)
        frame_seqlen = q.shape&#91;1&#93;
        logger.warning(&quot;grid_sizes not provided, using q.shape&#91;1&#93; as frame_seqlen&quot;)

    current_end = current_start + q.shape&#91;1&#93;
    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&#91;&quot;k&quot;&#93; = kv_cache&#91;&quot;k&quot;&#93;.detach().clone()
        kv_cache&#91;&quot;v&quot;&#93; = kv_cache&#91;&quot;v&quot;&#93;.detach().clone()
    else:
        kv_cache&#91;&quot;k&quot;&#93; = kv_cache&#91;&quot;k&quot;&#93;.detach()
        kv_cache&#91;&quot;v&quot;&#93; = kv_cache&#91;&quot;v&quot;&#93;.detach()

    kv_cache_size = kv_cache&#91;&quot;k&quot;&#93;.shape&#91;1&#93;
    num_new_tokens = q.shape&#91;1&#93;
    stored_key = k if relativistic else roped_key  # raw vs roped in cache
    global_end_index = (int(kv_cache&#91;&quot;global_end_index&quot;&#93;.item()) if isinstance(
        kv_cache&#91;&quot;global_end_index&quot;&#93;, torch.Tensor) else int(kv_cache&#91;&quot;global_end_index&quot;&#93;))
    local_end_index_prev = (int(kv_cache&#91;&quot;local_end_index&quot;&#93;.item()) if isinstance(
        kv_cache&#91;&quot;local_end_index&quot;&#93;, torch.Tensor) else int(kv_cache&#91;&quot;local_end_index&quot;&#93;))

    if (current_end &gt; global_end_index) and (num_new_tokens + local_end_index_prev &gt; 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&#91;&quot;k&quot;&#93;&#91;:, sink_tokens:sink_tokens + num_rolled_tokens&#93; = kv_cache&#91;&quot;k&quot;&#93;&#91;
            :,
            sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens,
        &#93;.clone()
        kv_cache&#91;&quot;v&quot;&#93;&#91;:, sink_tokens:sink_tokens + num_rolled_tokens&#93; = kv_cache&#91;&quot;v&quot;&#93;&#91;
            :,
            sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens,
        &#93;.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&#91;&quot;k&quot;&#93;&#91;:, local_start_index:local_end_index&#93; = stored_key
        kv_cache&#91;&quot;v&quot;&#93;&#91;:, local_start_index:local_end_index&#93; = 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&#91;&quot;k&quot;&#93;&#91;:, local_start_index:local_end_index&#93; = stored_key
        kv_cache&#91;&quot;v&quot;&#93;&#91;:, local_start_index:local_end_index&#93; = v

    k_for_attn, v_for_attn = retain_kv_with_sink(
        kv_cache&#91;&quot;k&quot;&#93;&#91;:, :local_end_index&#93;,
        kv_cache&#91;&quot;v&quot;&#93;&#91;:, :local_end_index&#93;,
        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&#91;1&#93;), int(grid_sizes&#91;2&#93;)
        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&#91;&quot;global_end_index&quot;&#93;, torch.Tensor):
        kv_cache&#91;&quot;global_end_index&quot;&#93;.fill_(current_end)
    else:
        kv_cache&#91;&quot;global_end_index&quot;&#93; = current_end
    if isinstance(kv_cache&#91;&quot;local_end_index&quot;&#93;, torch.Tensor):
        kv_cache&#91;&quot;local_end_index&quot;&#93;.fill_(local_end_index)
    else:
        kv_cache&#91;&quot;local_end_index&quot;&#93; = local_end_index

return x</code></pre></details>
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 = &quot;rms_norm_across_heads&quot;,
    cross_attn_norm: bool = False,
    eps: float = 1e-6,
    added_kv_proj_dim: int | None = None,
    supported_attention_backends: tuple&#91;AttentionBackendEnum, ...&#93;
    | None = None,
    prefix: str = &quot;&quot;,
    action_config: dict | None = None,
    block_idx: int = 0,
    rope_cache_policy: str = &quot;absolute&quot;,
):
    super().__init__()
    action_config = action_config or &#123;&#125;

    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 == &quot;rms_norm&quot;:
        self.norm_q = RMSNorm(dim_head, eps=eps)
        self.norm_k = RMSNorm(dim_head, eps=eps)
    elif qk_norm == &quot;rms_norm_across_heads&quot;:
        self.norm_q = RMSNorm(dim, eps=eps)
        self.norm_k = RMSNorm(dim, eps=eps)
    else:
        print(&quot;QK Norm type not supported&quot;)
        raise Exception
    assert cross_attn_norm is True
    self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
        dim,
        norm_type=&quot;layer&quot;,
        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=&quot;layer&quot;,
        eps=eps,
        elementwise_affine=False,
        dtype=torch.float32,
    )

    if len(action_config) &gt; 0 and block_idx in action_config.get(&quot;blocks&quot;, &#91;&#93;):
        self.action_model = ActionModule(
            heads_num=action_config&#91;&quot;heads_num&quot;&#93;,
            hidden_size=action_config&#91;&quot;hidden_size&quot;&#93;,
            img_hidden_size=action_config&#91;&quot;img_hidden_size&quot;&#93;,
            mouse_hidden_dim=action_config.get(&quot;mouse_hidden_dim&quot;, 0),
            keyboard_hidden_dim=action_config&#91;&quot;keyboard_hidden_dim&quot;&#93;,
            mouse_dim_in=action_config.get(&quot;mouse_dim_in&quot;, 0),
            keyboard_dim_in=action_config&#91;&quot;keyboard_dim_in&quot;&#93;,
            enable_mouse=action_config&#91;&quot;enable_mouse&quot;&#93;,
            enable_keyboard=action_config&#91;&quot;enable_keyboard&quot;&#93;,
            windows_size=action_config&#91;&quot;windows_size&quot;&#93;,
            rope_theta=action_config&#91;&quot;rope_theta&quot;&#93;,
            rope_dim_list=action_config&#91;&quot;rope_dim_list&quot;&#93;,
            mouse_qk_dim_list=action_config.get(&quot;mouse_qk_dim_list&quot;, &#91;8, 28, 28&#93;),
            patch_size=action_config&#91;&quot;patch_size&quot;&#93;,
            local_attn_size=local_attn_size,
            sink_size=sink_size,
            qk_norm=action_config&#91;&quot;qk_norm&quot;&#93;,
            qkv_bias=action_config&#91;&quot;qkv_bias&quot;&#93;,
            vae_time_compression_ratio=action_config&#91;&quot;vae_time_compression_ratio&quot;&#93;,
        )
    else:
        self.action_model = None

    self.ffn = MLP(dim, ffn_dim, act_type=&quot;gelu_pytorch_tanh&quot;)
    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&#91;torch.Tensor, torch.Tensor&#93;,
    block_mask: BlockMask,
    grid_sizes: tuple&#91;int, int, int&#93;,
    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,
) -&gt; torch.Tensor:
    if hidden_states.dim() == 4:
        hidden_states = hidden_states.squeeze(1)

    assert temb.ndim == 4
    num_frames = temb.shape&#91;1&#93;
    frame_seqlen = hidden_states.shape&#91;1&#93; // 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(&#91;0&#93;, 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&#91;0&#93;,
                grid_sizes&#91;1&#93;,
                grid_sizes&#91;2&#93;,
                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)</code></pre></details>
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&#91;1&#93;
frame_seqlen = hidden_states.shape&#91;1&#93; // 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(&#91;0&#93;, 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&#91;0&#93;,
            grid_sizes&#91;1&#93;,
            grid_sizes&#91;2&#93;,
            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)</code></pre></details>
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&#91;str, Any&#93;,
    **kwargs,
) -&gt; 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, &quot;arch_config&quot;, None)
    self.local_attn_size = (getattr(
        arch_cfg,
        &quot;local_attn_size&quot;,
        getattr(config, &quot;local_attn_size&quot;, -1),
    ) if arch_cfg else getattr(config, &quot;local_attn_size&quot;, -1))
    self.sink_size = (getattr(arch_cfg, &quot;sink_size&quot;, getattr(config, &quot;sink_size&quot;, 0)) if arch_cfg else getattr(
        config, &quot;sink_size&quot;, 0))
    if self.sink_size &lt; 0:
        raise ValueError(&quot;sink_size must be non-negative&quot;)
    if self.local_attn_size != -1 and self.sink_size &gt;= self.local_attn_size:
        raise ValueError(&quot;sink_size must be smaller than local_attn_size for &quot;
                         &quot;MatrixGame2 causal attention&quot;)
    self.rope_cache_policy = (getattr(arch_cfg, &quot;rope_cache_policy&quot;, getattr(
        config, &quot;rope_cache_policy&quot;, &quot;absolute&quot;)) if arch_cfg else getattr(config, &quot;rope_cache_policy&quot;, &quot;absolute&quot;))

    # 1. Patch &amp; 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, &quot;arch_config&quot;, None)
    self.action_config = (getattr(arch_cfg, &quot;action_config&quot;, &#123;&#125;) if arch_cfg is not None else &#123;&#125;)

    # 3. Transformer blocks
    self.blocks = nn.ModuleList(&#91;
        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&quot;&#123;getattr(config, &#x27;prefix&#x27;, &#x27;Wan&#x27;)&#125;.blocks.&#123;i&#125;&quot;,
            action_config=self.action_config,
            block_idx=i,
            rope_cache_policy=self.rope_cache_policy,
        ) for i in range(config.num_layers)
    &#93;)

    # 4. Output norm &amp; projection
    self.norm_out = LayerNormScaleShift(
        inner_dim,
        norm_type=&quot;layer&quot;,
        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,
        &quot;num_frames_per_block&quot;,
        getattr(config, &quot;num_frame_per_block&quot;, 1),
    ) if arch_cfg else getattr(config, &quot;num_frame_per_block&quot;, 1))

    self.__post_init__()

def set_causal_attention_window(
    self,
    *,
    local_attn_size: int | None = None,
    sink_size: int | None = None,
) -&gt; 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 &lt; 0:
        raise ValueError(&quot;sink_size must be non-negative&quot;)
    if local_attn_size != -1 and sink_size &gt;= local_attn_size:
        raise ValueError(&quot;sink_size must be smaller than local_attn_size for &quot;
                         &quot;MatrixGame2 causal attention&quot;)

    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, &quot;local_attn_size&quot;):
            block.local_attn_size = local_attn_size
        if hasattr(block, &quot;sink_size&quot;):
            block.sink_size = sink_size

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

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

def set_sink_size(self, sink_size: int) -&gt; 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,
) -&gt; 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&#91;tmp:tmp + frame_seqlen * num_frame_per_block&#93; = (tmp + frame_seqlen * num_frame_per_block)

    def attention_mask(b, h, q_idx, kv_idx):
        if local_attn_size == -1:
            return (kv_idx &lt; ends&#91;q_idx&#93;) | (q_idx == kv_idx)
        else:
            return ((kv_idx &lt; ends&#91;q_idx&#93;)
                    &amp; ((kv_idx &lt; sink_size * frame_seqlen)
                       | (kv_idx &gt;= (ends&#91;q_idx&#93; - 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(
            &quot;cache a block wise causal mask with block size of %s frames&quot;,
            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,
) -&gt; 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&#91;tmp:tmp + frame_seqlen * num_frame_per_block&#93; = cnt
        cnt += num_frame_per_block

    def attention_mask2(b, h, q_idx, kv_idx):
        if local_attn_size == -1:
            return (kv_idx &lt; ends2&#91;q_idx&#93;) | (q_idx == kv_idx)
        else:
            return ((kv_idx &lt; ends2&#91;q_idx&#93;)
                    &amp; ((kv_idx &lt; sink_size) | (kv_idx &gt;= (ends2&#91;q_idx&#93; - 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(
            &quot;cache a block wise causal mask for keyboard with block size of %s frames&quot;,
            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,
) -&gt; 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&#91;tmp:tmp + frame_seqlen * num_frame_per_block&#93; = cnt
        cnt += num_frame_per_block

    def attention_mask2(b, h, q_idx, kv_idx):
        if local_attn_size == -1:
            return (kv_idx &lt; ends2&#91;q_idx&#93;) | (q_idx == kv_idx)
        else:
            return ((kv_idx &lt; ends2&#91;q_idx&#93;)
                    &amp; ((kv_idx &lt; sink_size) | (kv_idx &gt;= (ends2&#91;q_idx&#93; - 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(
            &quot;cache a block wise causal mask for action with block size of %s frames&quot;,
            num_frame_per_block,
        )

    return block_mask2

def _forward_inference(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor | list&#91;torch.Tensor&#93;,
    timestep: torch.LongTensor,
    encoder_hidden_states_image: torch.Tensor
    | list&#91;torch.Tensor&#93;
    | 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,
) -&gt; torch.Tensor:
    # Extract num_frame_per_block from kwargs
    effective_num_frame_per_block = kwargs.pop(&quot;num_frame_per_block&quot;, 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) &gt; 0)

    if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states&#91;0&#93;
    if isinstance(encoder_hidden_states_image, list):
        encoder_hidden_states_image = (encoder_hidden_states_image&#91;0&#93;
                                       if len(encoder_hidden_states_image) &gt; 0 else None)

    ctx = get_forward_context()
    batch = getattr(ctx, &quot;forward_batch&quot;, None)

    if (batch is not None and getattr(batch, &quot;image_latent&quot;, 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&#91;2&#93; &gt;= end:
                cond = cond&#91;:, :, start:end&#93;
            elif cond.shape&#91;2&#93; &gt; start:
                cond = cond&#91;:, :, start:&#93;
                pad_frames = num_frames - cond.shape&#91;2&#93;
                if pad_frames &gt; 0:
                    pad = torch.zeros(
                        cond.shape&#91;0&#93;,
                        cond.shape&#91;1&#93;,
                        pad_frames,
                        cond.shape&#91;3&#93;,
                        cond.shape&#91;4&#93;,
                        device=cond.device,
                        dtype=cond.dtype,
                    )
                    cond = torch.cat(&#91;cond, pad&#93;, dim=2)
            else:
                pad = torch.zeros(
                    cond.shape&#91;0&#93;,
                    cond.shape&#91;1&#93;,
                    num_frames,
                    cond.shape&#91;3&#93;,
                    cond.shape&#91;4&#93;,
                    device=cond.device,
                    dtype=cond.dtype,
                )
                cond = pad
            hidden_states = torch.cat(&#91;hidden_states, cond.to(dtype=hidden_states.dtype)&#93;, 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 = &#91;d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)&#93;
    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&#91;0&#93;
        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(
                &#91;encoder_hidden_states_image, encoder_hidden_states_text&#93;,
                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 = &#91;None&#93; * len(self.blocks)
    if kv_cache_mouse is None:
        kv_cache_mouse = &#91;None&#93; * len(self.blocks)
    if kv_cache_keyboard is None:
        kv_cache_keyboard = &#91;None&#93; * len(self.blocks)
    if crossattn_cache is None:
        crossattn_cache = &#91;None&#93; * len(self.blocks)

    for block_index, block in enumerate(self.blocks):
        causal_kwargs = &#123;
            &quot;kv_cache&quot;: kv_cache&#91;block_index&#93;,
            &quot;kv_cache_mouse&quot;: kv_cache_mouse&#91;block_index&#93; if kv_cache_mouse else None,
            &quot;kv_cache_keyboard&quot;: kv_cache_keyboard&#91;block_index&#93; if kv_cache_keyboard else None,
            &quot;crossattn_cache&quot;: crossattn_cache&#91;block_index&#93; if crossattn_cache else None,
            &quot;current_start&quot;: current_start,
            &quot;cache_start&quot;: cache_start,
        &#125;
        if torch.is_grad_enabled() and self.gradient_checkpointing:
            causal_kwargs&#91;&quot;num_frame_per_block&quot;&#93; = 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&#91;torch.Tensor&#93;,
    timestep: torch.LongTensor,
    encoder_hidden_states_image: torch.Tensor
    | list&#91;torch.Tensor&#93;
    | None = None,
    mouse_cond: torch.Tensor | None = None,
    keyboard_cond: torch.Tensor | None = None,
    **kwargs,
) -&gt; 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) &gt; 0)

    if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states&#91;0&#93;
    if isinstance(encoder_hidden_states_image, list):
        encoder_hidden_states_image = (encoder_hidden_states_image&#91;0&#93;
                                       if len(encoder_hidden_states_image) &gt; 0 else None)

    ctx = get_forward_context()
    batch = getattr(ctx, &quot;forward_batch&quot;, None)
    if (batch is not None and getattr(batch, &quot;image_latent&quot;, 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&#91;2&#93; &gt;= num_frames:
                cond = cond&#91;:, :, :num_frames&#93;
            else:
                pad_frames = num_frames - cond.shape&#91;2&#93;
                pad = torch.zeros(
                    cond.shape&#91;0&#93;,
                    cond.shape&#91;1&#93;,
                    pad_frames,
                    cond.shape&#91;3&#93;,
                    cond.shape&#91;4&#93;,
                    device=cond.device,
                    dtype=cond.dtype,
                )
                cond = torch.cat(&#91;cond, pad&#93;, dim=2)
            hidden_states = torch.cat(&#91;hidden_states, cond.to(dtype=hidden_states.dtype)&#93;, 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 = &#91;d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)&#93;
    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&#91;0&#93;
        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(
                &#91;encoder_hidden_states_image, encoder_hidden_states_text&#93;,
                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(&quot;kv_cache&quot;, &#91;None&#93; * len(self.blocks))
    kv_cache_mouse = kwargs.pop(&quot;kv_cache_mouse&quot;, &#91;None&#93; * len(self.blocks))
    kv_cache_keyboard = kwargs.pop(&quot;kv_cache_keyboard&quot;, &#91;None&#93; * len(self.blocks))
    crossattn_cache = kwargs.pop(&quot;crossattn_cache&quot;, &#91;None&#93; * len(self.blocks))

    for i, block in enumerate(self.blocks):
        block_crossattn_cache = None
        if crossattn_cache is not None:
            if crossattn_cache&#91;i&#93; is None:
                crossattn_cache&#91;i&#93; = &#123;&quot;is_init&quot;: False&#125;
            block_crossattn_cache = crossattn_cache&#91;i&#93;

        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&#91;i&#93;,
                kv_cache_mouse=kv_cache_mouse&#91;i&#93;,
                kv_cache_keyboard=kv_cache_keyboard&#91;i&#93;,
                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&#91;i&#93;,
                kv_cache_mouse=kv_cache_mouse&#91;i&#93;,
                kv_cache_keyboard=kv_cache_keyboard&#91;i&#93;,
                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(&quot;kv_cache&quot;) 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&#91;int, int, int&#93;) -&gt; 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&#91;:, :f * h * w&#93;.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</code></pre></details>
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, &quot;local_attn_size&quot;):
        block.local_attn_size = local_attn_size
    if hasattr(block, &quot;sink_size&quot;):
        block.sink_size = sink_size

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

    action_model = getattr(block, &quot;action_model&quot;, None)
    if action_model is not None:
        if hasattr(action_model, &quot;local_attn_size&quot;):
            action_model.local_attn_size = local_attn_size
        if hasattr(action_model, &quot;sink_size&quot;):
            action_model.sink_size = sink_size</code></pre></details>
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&#91;:, :f * h * w&#93;.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</code></pre></details>
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</code></pre></details>