diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index 385b2a61a..b3b3593ed 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -16,8 +16,8 @@ from fastvideo.layers.rotary_embedding import ( ) from fastvideo.platforms import AttentionBackendEnum from .kv_cache import ( - get_attended_kv_without_update, - update_kv_cache_and_get_attended_kv, + attend_with_kv_cache, + KVCache, ) @@ -99,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, @@ -115,30 +115,18 @@ def _update_kv_cache_and_attend( k.shape[1] if use_k_for_num_tokens else q.shape[1] ) == num_frame_per_block num_new_tokens = k.shape[1] if use_k_for_num_tokens else q.shape[1] - if update_kv_cache: - cached_k, cached_v = update_kv_cache_and_get_attended_kv( - kv_cache=kv_cache, - new_k=k, - new_v=v, - num_new_tokens=num_new_tokens, - max_attn_size=local_attn_size, - store_first_only=store_first_only, - ) - else: - cached_k, cached_v = get_attended_kv_without_update( - kv_cache=kv_cache, - new_k=k, - new_v=v, - num_new_tokens=num_new_tokens, - max_attn_size=local_attn_size, - store_first_only=store_first_only, - ) - - 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) - - return attn_layer(q, cached_k, cached_v) + 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, + ) class ActionModule(nn.Module): @@ -371,7 +359,7 @@ 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, @@ -489,7 +477,7 @@ 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, @@ -648,8 +636,8 @@ 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, diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 4f4e7f751..9a2d2320c 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -45,8 +45,8 @@ from fastvideo.platforms import AttentionBackendEnum, current_platform from .action_module import ActionModule from .kv_cache import ( - get_attended_kv_without_update, - update_kv_cache_and_get_attended_kv, + attend_with_kv_cache, + KVCache, ) from .model import MatrixGameCrossAttention @@ -181,7 +181,7 @@ 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, @@ -277,28 +277,20 @@ class CausalMatrixGameSelfAttention(nn.Module): if self.local_attn_size == -1 else self.local_attn_size * frame_seqlen ) - if update_kv_cache: - k_for_attn, v_for_attn = update_kv_cache_and_get_attended_kv( - kv_cache=kv_cache, - new_k=roped_key, - new_v=v, - num_new_tokens=roped_query.shape[1], - max_attn_size=max_attention_size, - ) - else: - k_for_attn, v_for_attn = get_attended_kv_without_update( - kv_cache=kv_cache, - new_k=roped_key, - new_v=v, - num_new_tokens=roped_query.shape[1], - max_attn_size=max_attention_size, - ) - - 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) + 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], + ) return x @@ -423,9 +415,9 @@ 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, @@ -827,10 +819,10 @@ 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, diff --git a/fastvideo/models/dits/matrixgame/kv_cache.py b/fastvideo/models/dits/matrixgame/kv_cache.py index 8a78d88da..be05c5af6 100644 --- a/fastvideo/models/dits/matrixgame/kv_cache.py +++ b/fastvideo/models/dits/matrixgame/kv_cache.py @@ -1,8 +1,81 @@ # 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 length_int(self) -> int: + return int(self.length.item()) + + def set_length(self, value: int) -> None: + self.length.fill_(int(value)) + + +@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 _store_first_tokens( new_tensor: torch.Tensor, num_new_tokens: int, @@ -14,20 +87,16 @@ def _store_first_tokens( def get_attended_kv_without_update( - kv_cache: dict[str, torch.Tensor | int], + kv_cache: KVCache, new_k: torch.Tensor, new_v: torch.Tensor, num_new_tokens: int, max_attn_size: int, store_first_only: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - kv_capacity = kv_cache["k"].shape[1] - target_batch = kv_cache["k"].shape[0] - cache_length = ( - int(kv_cache["length"].item()) - if isinstance(kv_cache["length"], torch.Tensor) - else int(kv_cache["length"]) - ) + kv_capacity = kv_cache.capacity + target_batch = kv_cache.batch_size + cache_length = kv_cache.length_int() if store_first_only: write_k = _store_first_tokens(new_k, num_new_tokens, target_batch) @@ -36,8 +105,8 @@ def get_attended_kv_without_update( write_k = new_k write_v = new_v - cache_k = kv_cache["k"][:, :cache_length] - cache_v = kv_cache["v"][:, :cache_length] + cache_k = kv_cache.k[:, :cache_length] + cache_v = kv_cache.v[:, :cache_length] if num_new_tokens + cache_length > kv_capacity: num_evicted_tokens = num_new_tokens + cache_length - kv_capacity @@ -57,26 +126,22 @@ def get_attended_kv_without_update( def update_kv_cache_and_get_attended_kv( - kv_cache: dict[str, torch.Tensor | int], + kv_cache: KVCache, new_k: torch.Tensor, new_v: torch.Tensor, num_new_tokens: int, max_attn_size: int, store_first_only: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - kv_capacity = kv_cache["k"].shape[1] - cache_length = ( - int(kv_cache["length"].item()) - if isinstance(kv_cache["length"], torch.Tensor) - else int(kv_cache["length"]) - ) + kv_capacity = kv_cache.capacity + cache_length = kv_cache.length_int() - cache_k = kv_cache["k"][:, :cache_length] - cache_v = kv_cache["v"][:, :cache_length] + cache_k = kv_cache.k[:, :cache_length] + cache_v = kv_cache.v[:, :cache_length] if store_first_only: - write_k = _store_first_tokens(new_k, num_new_tokens, kv_cache["k"].shape[0]) - write_v = _store_first_tokens(new_v, num_new_tokens, kv_cache["v"].shape[0]) + write_k = _store_first_tokens(new_k, num_new_tokens, kv_cache.batch_size) + write_v = _store_first_tokens(new_v, num_new_tokens, kv_cache.batch_size) else: write_k = new_k write_v = new_v @@ -92,16 +157,55 @@ def update_kv_cache_and_get_attended_kv( updated_v = torch.cat([cache_v, write_v], dim=1) updated_length = min(updated_k.shape[1], kv_capacity) - kv_cache["k"][:, :updated_length] = updated_k[:, :updated_length] - kv_cache["v"][:, :updated_length] = updated_v[:, :updated_length] + kv_cache.k[:, :updated_length] = updated_k[:, :updated_length] + kv_cache.v[:, :updated_length] = updated_v[:, :updated_length] attn_start = max(0, updated_length - max_attn_size) - cached_k = kv_cache["k"][:, attn_start:updated_length] - cached_v = kv_cache["v"][:, attn_start:updated_length] + cached_k = kv_cache.k[:, attn_start:updated_length] + cached_v = kv_cache.v[:, attn_start:updated_length] - if isinstance(kv_cache["length"], torch.Tensor): - kv_cache["length"].fill_(updated_length) - else: - kv_cache["length"] = updated_length + kv_cache.set_length(updated_length) return cached_k, cached_v + + +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""" + tokens_to_add = q.shape[1] if num_new_tokens is None else num_new_tokens + + if update_kv_cache: + cached_k, cached_v = update_kv_cache_and_get_attended_kv( + kv_cache=kv_cache, + new_k=new_k, + new_v=new_v, + num_new_tokens=tokens_to_add, + max_attn_size=max_attn_size, + store_first_only=store_first_only, + ) + else: + cached_k, cached_v = get_attended_kv_without_update( + kv_cache=kv_cache, + new_k=new_k, + new_v=new_v, + num_new_tokens=tokens_to_add, + max_attn_size=max_attn_size, + store_first_only=store_first_only, + ) + + 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) + + 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 39d74d19c..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,24 +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), - "length": - 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 @@ -346,40 +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), - "length": - 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), - "length": - 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