Merge branch 'feat/kaiqin/mg-kvcache-refactor' into feat/kaiqin/mg_solaris

fastvideo/models/dits/matrixgame/action_module.py

@@ -15,6 +15,10 @@ from fastvideo.layers.rotary_embedding import (
     _apply_rotary_emb,
 )
 from fastvideo.platforms import AttentionBackendEnum
+from .kv_cache import (
+    attend_with_kv_cache,
+    KVCache,
+)
 
 
 DISABLE_COMPILE = False
@@ -95,7 +99,7 @@ def _update_kv_cache_and_attend(
     q: torch.Tensor,
     k: torch.Tensor,
     v: torch.Tensor,
-    kv_cache: dict[str, torch.Tensor | int],
+    kv_cache: KVCache,
     attn_layer: LocalAttention,
     start_frame: int,
     num_frame_per_block: int,
@@ -104,129 +108,25 @@ def _update_kv_cache_and_attend(
     use_k_for_num_tokens: bool = False,
     store_first_only: bool = False,
     repeat_factor: int | None = None,
+    update_kv_cache: bool = True,
 ) -> torch.Tensor:
-    """
-    Update KV cache with new tokens and perform attention with cached values.
-
-    Args:
-        q: Query tensor
-        k: Key tensor
-        v: Value tensor
-        kv_cache: Dictionary containing cache tensors and indices
-        attn_layer: Attention layer to use
-        start_frame: Starting frame index
-        num_frame_per_block: Number of frames per block
-        local_attn_size: Maximum attention window size
-        use_k_for_num_tokens: If True, use k.shape[1] for num_new_tokens, else use q.shape[1]
-        store_first_only: If True, only store k[:1] and v[:1] in cache (for keyboard with rope)
-        repeat_factor: If provided, repeat cached k,v by this factor when retrieving (for keyboard with rope)
-
-    Returns:
-        Attention output tensor
-    """
-    current_start = start_frame
-    current_end = current_start + (
-        k.shape[1] if use_k_for_num_tokens else q.shape[1]
-    )
-
+    """Update KV cache."""
     assert (
         k.shape[1] if use_k_for_num_tokens else q.shape[1]
     ) == num_frame_per_block
-
-    sink_size = 0
-    max_attention_size = local_attn_size
-    sink_tokens = sink_size * 1
-    kv_cache_size = kv_cache["k"].shape[1]
     num_new_tokens = k.shape[1] if use_k_for_num_tokens else q.shape[1]
-
-    original_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"])
+    return attend_with_kv_cache(
+        q=q,
+        new_k=k,
+        new_v=v,
+        kv_cache=kv_cache,
+        max_attn_size=local_attn_size,
+        attend_fn=attn_layer,
+        update_kv_cache=update_kv_cache,
+        store_first_only=store_first_only,
+        repeat_factor=repeat_factor,
+        num_new_tokens=num_new_tokens,
     )
-    original_local_end_index = (
-        int(kv_cache["local_end_index"].item())
-        if isinstance(kv_cache["local_end_index"], torch.Tensor)
-        else int(kv_cache["local_end_index"])
-    )
-
-    # Check if we need to evict tokens
-    if (current_end > original_global_end_index) and (
-        num_new_tokens + original_local_end_index > kv_cache_size
-    ):
-        num_evicted_tokens = (
-            num_new_tokens + original_local_end_index - kv_cache_size
-        )
-        num_rolled_tokens = (
-            original_local_end_index
-            - num_evicted_tokens
-            - sink_tokens
-        )
-        # Roll k cache
-        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()
-        )
-        # Roll v cache
-        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()
-        )
-        # Calculate indices with eviction adjustment
-        local_end_index = (
-            original_local_end_index
-            + current_end
-            - original_global_end_index
-            - num_evicted_tokens
-        )
-        local_start_index = local_end_index - num_new_tokens
-    else:
-        # Calculate indices without eviction
-        local_end_index = (
-            original_local_end_index
-            + current_end
-            - original_global_end_index
-        )
-        local_start_index = local_end_index - num_new_tokens
-
-    # Store new k, v in cache
-    if store_first_only:
-        kv_cache["k"][:, local_start_index:local_end_index] = k[:1]
-        kv_cache["v"][:, local_start_index:local_end_index] = v[:1]
-    else:
-        kv_cache["k"][:, local_start_index:local_end_index] = k
-        kv_cache["v"][:, local_start_index:local_end_index] = v
-
-    # Retrieve from cache and perform attention
-    cache_start = max(0, local_end_index - max_attention_size)
-    cached_k = kv_cache["k"][:, cache_start:local_end_index]
-    cached_v = kv_cache["v"][:, cache_start:local_end_index]
-
-    if repeat_factor is not None:
-        cached_k = cached_k.repeat(repeat_factor, 1, 1, 1)
-        cached_v = cached_v.repeat(repeat_factor, 1, 1, 1)
-
-    attn = attn_layer(q, cached_k, cached_v)
-
-    # Update indices
-    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 attn
 
 
 class ActionModule(nn.Module):
@@ -459,11 +359,12 @@ class ActionModule(nn.Module):
         mouse_condition: torch.Tensor,
         *,
         is_causal: bool,
-        kv_cache_mouse: dict[str, torch.Tensor] | None,
+        kv_cache_mouse: KVCache | None,
         pad_t: int,
         num_frame_per_block: int,
         block_mask_mouse: BlockMask | None,
         start_frame: int,
+        update_kv_cache: bool,
         freqs_cis: tuple[torch.Tensor, torch.Tensor],
         N_feats: int,
         B: int,
@@ -559,6 +460,7 @@ class ActionModule(nn.Module):
                     num_frame_per_block,
                     self.local_attn_size,
                     use_k_for_num_tokens=False,
+                    update_kv_cache=update_kv_cache,
                 )
         else:
             attn = self.mouse_attn_layer(q, k, v)
@@ -575,11 +477,12 @@ class ActionModule(nn.Module):
         *,
         is_causal: bool,
         use_rope_keyboard: bool,
-        kv_cache_keyboard: dict[str, torch.Tensor] | None,
+        kv_cache_keyboard: KVCache | None,
         pad_t: int,
         num_frame_per_block: int,
         block_mask_keyboard: BlockMask | None,
         start_frame: int,
+        update_kv_cache: bool,
         freqs_cis: tuple[torch.Tensor, torch.Tensor],
         N_feats: int,
         B: int,
@@ -685,6 +588,7 @@ class ActionModule(nn.Module):
                         use_k_for_num_tokens=True,
                         store_first_only=True,
                         repeat_factor=S,
+                        update_kv_cache=update_kv_cache,
                     )
             else:
                 attn = self.keyboard_attn_layer(q, k, v)
@@ -713,6 +617,7 @@ class ActionModule(nn.Module):
                         num_frame_per_block,
                         self.local_attn_size,
                         use_k_for_num_tokens=True,
+                        update_kv_cache=update_kv_cache,
                     )
             else:
                 attn = self.keyboard_attn_layer(q, k, v)
@@ -731,11 +636,12 @@ class ActionModule(nn.Module):
         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,
+        kv_cache_mouse: KVCache | None = None,
+        kv_cache_keyboard: KVCache | None = None,
         start_frame: int = 0,
         use_rope_keyboard: bool = True,
         num_frame_per_block: int = 3,
+        update_kv_cache: bool = True,
     ):
         """
         hidden_states: B, tt*th*tw, C
@@ -805,6 +711,7 @@ class ActionModule(nn.Module):
                 num_frame_per_block=num_frame_per_block,
                 block_mask_mouse=block_mask_mouse,
                 start_frame=start_frame,
+                update_kv_cache=update_kv_cache,
                 freqs_cis=freqs_cis,
                 N_feats=N_feats,
                 B=B,
@@ -828,6 +735,7 @@ class ActionModule(nn.Module):
                 num_frame_per_block=num_frame_per_block,
                 block_mask_keyboard=block_mask_keyboard,
                 start_frame=start_frame,
+                update_kv_cache=update_kv_cache,
                 freqs_cis=freqs_cis,
                 N_feats=N_feats,
                 B=B,

fastvideo/models/dits/matrixgame/causal_model.py

@@ -44,6 +44,10 @@ from fastvideo.models.dits.wanvideo import (
 from fastvideo.platforms import AttentionBackendEnum, current_platform
 
 from .action_module import ActionModule
+from .kv_cache import (
+    attend_with_kv_cache,
+    KVCache,
+)
 from .model import MatrixGameCrossAttention
 
 logger = init_logger(__name__)
@@ -144,7 +148,6 @@ class CausalMatrixGameSelfAttention(nn.Module):
         dim: int,
         num_heads: int,
         local_attn_size: int = -1,
-        sink_size: int = 0,
         qk_norm: bool = True,
         eps: float = 1e-6,
     ) -> None:
@@ -154,7 +157,6 @@ class CausalMatrixGameSelfAttention(nn.Module):
         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._freqs_cache = None
@@ -179,9 +181,10 @@ class CausalMatrixGameSelfAttention(nn.Module):
         freqs_cis: tuple[torch.Tensor, torch.Tensor],
         block_mask: BlockMask,
         grid_sizes: tuple[int, int, int],
-        kv_cache: dict | None = None,
+        kv_cache: KVCache | None = None,
         current_start: int = 0,
         cache_start: int | None = None,
+        update_kv_cache: bool = True,
     ):
         if cache_start is None:
             cache_start = current_start
@@ -268,87 +271,26 @@ class CausalMatrixGameSelfAttention(nn.Module):
                     "grid_sizes not provided, using q.shape[1] as frame_seqlen"
                 )
 
-            current_end = current_start + roped_query.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
             )
-
-            kv_cache_size = kv_cache["k"].shape[1]
-            num_new_tokens = roped_query.shape[1]
-            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"])
+            x = attend_with_kv_cache(
+                q=roped_query,
+                new_k=roped_key,
+                new_v=v,
+                kv_cache=kv_cache,
+                max_attn_size=max_attention_size,
+                attend_fn=lambda q_t, k_t, v_t: torch.nn.functional.scaled_dot_product_attention(
+                    q_t.transpose(1, 2),
+                    k_t.transpose(1, 2),
+                    v_t.transpose(1, 2),
+                ).transpose(1, 2),
+                update_kv_cache=update_kv_cache,
+                num_new_tokens=roped_query.shape[1],
             )
-            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] = roped_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] = roped_key
-                kv_cache["v"][:, local_start_index:local_end_index] = v
-
-            kv_start = max(0, local_end_index - max_attention_size)
-            k_for_attn = kv_cache["k"][:, kv_start:local_end_index]
-            v_for_attn = kv_cache["v"][:, kv_start:local_end_index]
-
-            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
 
@@ -360,7 +302,6 @@ class CausalMatrixGameTransformerBlock(nn.Module):
         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,
@@ -384,7 +325,6 @@ class CausalMatrixGameTransformerBlock(nn.Module):
             dim,
             num_heads,
             local_attn_size=local_attn_size,
-            sink_size=sink_size,
             qk_norm=qk_norm,
             eps=eps,
         )
@@ -475,12 +415,13 @@ class CausalMatrixGameTransformerBlock(nn.Module):
         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,
+        kv_cache: KVCache | None = None,
+        kv_cache_mouse: KVCache | None = None,
+        kv_cache_keyboard: KVCache | None = None,
         crossattn_cache: dict | None = None,
         current_start: int = 0,
         cache_start: int | None = None,
+        update_kv_cache: bool = True,
     ) -> torch.Tensor:
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
@@ -528,6 +469,7 @@ class CausalMatrixGameTransformerBlock(nn.Module):
             kv_cache,
             current_start,
             cache_start,
+            update_kv_cache=update_kv_cache,
         )
         attn_output = attn_output.flatten(2)
         attn_output, _ = self.to_out(attn_output)
@@ -567,6 +509,7 @@ class CausalMatrixGameTransformerBlock(nn.Module):
                     start_frame=start_frame,
                     use_rope_keyboard=use_rope_keyboard,
                     num_frame_per_block=num_frame_per_block,
+                    update_kv_cache=update_kv_cache,
                 )
 
         ff_output = self.ffn(
@@ -630,11 +573,6 @@ class CausalMatrixGameWanModel(BaseDiT):
             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)
-        )
 
         # 1. Patch & position embedding
         self.patch_embedding = PatchEmbed(
@@ -667,7 +605,6 @@ class CausalMatrixGameWanModel(BaseDiT):
                     config.ffn_dim,
                     config.num_attention_heads,
                     self.local_attn_size,
-                    self.sink_size,
                     config.qk_norm,
                     config.cross_attn_norm,
                     config.eps,
@@ -882,13 +819,14 @@ class CausalMatrixGameWanModel(BaseDiT):
         | 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,
+        kv_cache: list[KVCache] | None = None,
+        kv_cache_mouse: list[KVCache] | None = None,
+        kv_cache_keyboard: list[KVCache] | None = None,
+        crossattn_cache: list[dict] | None = None,
         current_start: int = 0,
         cache_start: int = 0,
         start_frame: int = 0,
+        update_kv_cache: bool = True,
         **kwargs,
     ) -> torch.Tensor:
         # Extract num_frame_per_block from kwargs
@@ -1089,6 +1027,7 @@ class CausalMatrixGameWanModel(BaseDiT):
                         "current_start": current_start,
                         "cache_start": cache_start,
                         "num_frame_per_block": effective_num_frame_per_block,
+                        "update_kv_cache": update_kv_cache,
                     }
                 )
                 hidden_states = self._gradient_checkpointing_func(
@@ -1109,6 +1048,7 @@ class CausalMatrixGameWanModel(BaseDiT):
                         else None,
                         "current_start": current_start,
                         "cache_start": cache_start,
+                        "update_kv_cache": update_kv_cache,
                     }
                 )
                 hidden_states = block(

fastvideo/models/dits/matrixgame/kv_cache.py

@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class KVCache:
+
+    k: torch.Tensor
+    v: torch.Tensor
+    length: torch.Tensor
+
+    @classmethod
+    def zeros(
+        cls,
+        batch_size: int,
+        cache_size: int,
+        num_heads: int,
+        head_dim: int,
+        *,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> "KVCache":
+        return cls(
+            k=torch.zeros(
+                [batch_size, cache_size, num_heads, head_dim],
+                dtype=dtype,
+                device=device,
+            ),
+            v=torch.zeros(
+                [batch_size, cache_size, num_heads, head_dim],
+                dtype=dtype,
+                device=device,
+            ),
+            length=torch.tensor([0], dtype=torch.long, device=device),
+        )
+
+    @classmethod
+    def from_dict(cls, kv_cache: dict[str, torch.Tensor | int]) -> "KVCache":
+        length = kv_cache["length"]
+        if not isinstance(length, torch.Tensor):
+            length = torch.tensor([int(length)], dtype=torch.long, device=kv_cache["k"].device)
+        return cls(k=kv_cache["k"], v=kv_cache["v"], length=length)
+
+    def to_dict(self) -> dict[str, torch.Tensor]:
+        return {"k": self.k, "v": self.v, "length": self.length}
+
+    @property
+    def capacity(self) -> int:
+        return self.k.shape[1]
+
+    @property
+    def batch_size(self) -> int:
+        return self.k.shape[0]
+
+    def get_len(self) -> int:
+        return int(self.length.item())
+
+    def _preprocess_kv(self, new_k: torch.Tensor, new_v: torch.Tensor, num_new_tokens: int, store_first_only: bool):
+        target_batch = self.batch_size
+        if store_first_only:
+            # Expand first token to match required sequence length
+            write_k = new_k[:target_batch, :1].expand(-1, num_new_tokens, -1, -1)
+            write_v = new_v[:target_batch, :1].expand(-1, num_new_tokens, -1, -1)
+        else:
+            write_k = new_k[:target_batch]
+            write_v = new_v[:target_batch]
+        return write_k, write_v
+
+    def update(
+        self,
+        new_k: torch.Tensor,
+        new_v: torch.Tensor,
+        max_attn_size: int,
+        store_first_only: bool = False,
+        repeat_factor: int | None = None,
+        num_new_tokens: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        num_tokens = num_new_tokens if num_new_tokens is not None else new_k.shape[1]
+        write_k, write_v = self._preprocess_kv(new_k, new_v, num_tokens, store_first_only)
+        
+        cur_len = self.get_len()
+        
+        # Combine existing valid tokens with new ones
+        combined_k = torch.cat([self.k[:, :cur_len], write_k], dim=1)
+        combined_v = torch.cat([self.v[:, :cur_len], write_v], dim=1)
+        
+        # Keep only up to capacity
+        new_len = min(combined_k.shape[1], self.capacity)
+        self.k[:, :new_len] = combined_k[:, -new_len:]
+        self.v[:, :new_len] = combined_v[:, -new_len:]
+        self.length.fill_(new_len)
+        
+        return self._fetch_kv_slice(self.k, self.v, new_len, max_attn_size, repeat_factor)
+
+    def get_view(
+        self,
+        new_k: torch.Tensor,
+        new_v: torch.Tensor,
+        max_attn_size: int,
+        store_first_only: bool = False,
+        repeat_factor: int | None = None,
+        num_new_tokens: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        num_tokens = num_new_tokens if num_new_tokens is not None else new_k.shape[1]
+        write_k, write_v = self._preprocess_kv(new_k, new_v, num_tokens, store_first_only)
+        
+        cur_len = self.get_len()
+        combined_k = torch.cat([self.k[:, :cur_len], write_k], dim=1)
+        combined_v = torch.cat([self.v[:, :cur_len], write_v], dim=1)
+        
+        return self._fetch_kv_slice(combined_k, combined_v, combined_k.shape[1], max_attn_size, repeat_factor)
+
+    def _fetch_kv_slice(self, k, v, current_len, max_attn_size, repeat_factor):
+        start = max(0, current_len - max_attn_size)
+        out_k = k[:, start:current_len]
+        out_v = v[:, start:current_len]
+        
+        if repeat_factor is not None:
+            out_k = out_k.repeat(repeat_factor, 1, 1, 1)
+            out_v = out_v.repeat(repeat_factor, 1, 1, 1)
+        return out_k, out_v
+
+
+@dataclass
+class KVCacheDict:
+    kv_cache: KVCache
+    kv_cache_mouse: KVCache
+    kv_cache_keyboard: KVCache
+
+    def to_dict(self) -> dict[str, dict[str, torch.Tensor]]:
+        return {
+            "kv_cache": self.kv_cache.to_dict(),
+            "kv_cache_mouse": self.kv_cache_mouse.to_dict(),
+            "kv_cache_keyboard": self.kv_cache_keyboard.to_dict(),
+        }
+
+
+def attend_with_kv_cache(
+    q: torch.Tensor,
+    new_k: torch.Tensor,
+    new_v: torch.Tensor,
+    kv_cache: KVCache,
+    *,
+    max_attn_size: int,
+    attend_fn,
+    update_kv_cache: bool = True,
+    store_first_only: bool = False,
+    repeat_factor: int | None = None,
+    num_new_tokens: int | None = None,
+) -> torch.Tensor:
+    """Unified KV-cache attention path"""
+
+    if update_kv_cache:
+        cached_k, cached_v = kv_cache.update(
+            new_k, new_v, max_attn_size, store_first_only, repeat_factor, num_new_tokens
+        )
+    else:
+        cached_k, cached_v = kv_cache.get_view(
+            new_k, new_v, max_attn_size, store_first_only, repeat_factor, num_new_tokens
+        )
+    return attend_fn(q, cached_k, cached_v)

fastvideo/pipelines/stages/matrixgame_denoising.py

@@ -23,6 +23,8 @@ except ImportError:
     vsa_available = False
     VideoSparseAttentionBackend = None  # type: ignore
 
+from fastvideo.models.dits.matrixgame.kv_cache import KVCache
+
 logger = init_logger(__name__)
 
 
@@ -35,10 +37,10 @@ class BlockProcessingContext:
     block_idx: int
     start_index: int
 
-    kv_cache1: list[dict[Any, Any]]
-    kv_cache2: list[dict[Any, Any]] | None
-    kv_cache_mouse: list[dict[Any, Any]] | None
-    kv_cache_keyboard: list[dict[Any, Any]] | None
+    kv_cache1: list[KVCache]
+    kv_cache2: list[KVCache] | None
+    kv_cache_mouse: list[KVCache] | None
+    kv_cache_keyboard: list[KVCache] | None
     crossattn_cache: list[dict[Any, Any]]
 
     timesteps: torch.Tensor
@@ -55,7 +57,7 @@ class BlockProcessingContext:
     image_kwargs: dict[str, Any]
     pos_cond_kwargs: dict[str, Any]
 
-    def get_kv_cache(self, timestep_val: float) -> list[dict[Any, Any]]:
+    def get_kv_cache(self, timestep_val: float) -> list[KVCache]:
         if self.boundary_timestep is not None:
             if timestep_val >= self.boundary_timestep:
                 return self.kv_cache1
@@ -296,7 +298,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
         return action_kwargs
 
     def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype,
-                             device: torch.device) -> list[dict]:
+                             device: torch.device) -> list[KVCache]:
         kv_cache = []
         num_attention_heads = self.transformer.num_attention_heads
         attention_head_dim = getattr(
@@ -308,26 +310,15 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
             kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames
 
         for _ in range(self.num_transformer_blocks):
-            kv_cache.append({
-                "k":
-                torch.zeros([
-                    batch_size, kv_cache_size, num_attention_heads,
-                    attention_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "v":
-                torch.zeros([
-                    batch_size, kv_cache_size, num_attention_heads,
-                    attention_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "global_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-                "local_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-            })
+            kv_cache.append(
+                KVCache.zeros(
+                    batch_size=batch_size,
+                    cache_size=kv_cache_size,
+                    num_heads=num_attention_heads,
+                    head_dim=attention_head_dim,
+                    dtype=dtype,
+                    device=device,
+                ))
 
         return kv_cache
 
@@ -348,44 +339,24 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
             kv_cache_size = 15
 
         for _ in range(self.num_transformer_blocks):
-            kv_cache_keyboard.append({
-                "k":
-                torch.zeros([
-                    batch_size, kv_cache_size, action_heads, keyboard_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "v":
-                torch.zeros([
-                    batch_size, kv_cache_size, action_heads, keyboard_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "global_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-                "local_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-            })
-            kv_cache_mouse.append({
-                "k":
-                torch.zeros([
-                    batch_size * self.frame_seq_length, kv_cache_size,
-                    action_heads, mouse_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "v":
-                torch.zeros([
-                    batch_size * self.frame_seq_length, kv_cache_size,
-                    action_heads, mouse_head_dim
-                ],
-                            dtype=dtype,
-                            device=device),
-                "global_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-                "local_end_index":
-                torch.tensor([0], dtype=torch.long, device=device),
-            })
+            kv_cache_keyboard.append(
+                KVCache.zeros(
+                    batch_size=batch_size,
+                    cache_size=kv_cache_size,
+                    num_heads=action_heads,
+                    head_dim=keyboard_head_dim,
+                    dtype=dtype,
+                    device=device,
+                ))
+            kv_cache_mouse.append(
+                KVCache.zeros(
+                    batch_size=batch_size * self.frame_seq_length,
+                    cache_size=kv_cache_size,
+                    num_heads=action_heads,
+                    head_dim=mouse_head_dim,
+                    dtype=dtype,
+                    device=device,
+                ))
 
         return kv_cache_mouse, kv_cache_keyboard
 
@@ -495,6 +466,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
                     "crossattn_cache": ctx.crossattn_cache,
                     "current_start": start_index * self.frame_seq_length,
                     "start_frame": start_index,
+                    "update_kv_cache": False,
                 }
 
                 if self.use_action_module and current_model == self.transformer:
@@ -605,6 +577,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
                 "crossattn_cache": ctx.crossattn_cache,
                 "current_start": start_index * self.frame_seq_length,
                 "start_frame": start_index,
+                "update_kv_cache": True,
             }
 
             if self.use_action_module:
@@ -625,6 +598,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
                     crossattn_cache=ctx.crossattn_cache,
                     current_start=start_index * self.frame_seq_length,
                     start_frame=start_index,
+                    update_kv_cache=True,
                     **ctx.image_kwargs,
                     **ctx.pos_cond_kwargs,
                 )