diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index d2ea11925..b3b3593ed 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/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, diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 99bb8e134..9a2d2320c 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ 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( diff --git a/fastvideo/models/dits/matrixgame/kv_cache.py b/fastvideo/models/dits/matrixgame/kv_cache.py new file mode 100644 index 000000000..6d9f9f67b --- /dev/null +++ b/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) diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index cd9599bbb..f187f127a 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/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, )