From b32aa9cb0b264a85a60d75c4f5fff8d15595ad72 Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Fri, 27 Feb 2026 04:13:37 +0000 Subject: [PATCH 1/6] remove sinksize --- .../models/dits/matrixgame/action_module.py | 30 +++++-------------- .../models/dits/matrixgame/causal_model.py | 29 ++++-------------- 2 files changed, 12 insertions(+), 47 deletions(-) diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index d2ea11925..3a7b0ff53 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -133,9 +133,7 @@ def _update_kv_cache_and_attend( 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] @@ -157,29 +155,15 @@ def _update_kv_cache_and_attend( 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 - ) + num_rolled_tokens = original_local_end_index - num_evicted_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() - ) + kv_cache["k"][:, :num_rolled_tokens] = kv_cache["k"][ + :, num_evicted_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() - ) + kv_cache["v"][:, :num_rolled_tokens] = kv_cache["v"][ + :, num_evicted_tokens : num_evicted_tokens + num_rolled_tokens + ].clone() # Calculate indices with eviction adjustment local_end_index = ( original_local_end_index diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 99bb8e134..0627b96a6 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -144,7 +144,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 +153,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 @@ -269,7 +267,6 @@ class CausalMatrixGameSelfAttention(nn.Module): ) 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 = ( @@ -297,23 +294,15 @@ class CausalMatrixGameSelfAttention(nn.Module): 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"][ + num_rolled_tokens = local_end_index_prev - num_evicted_tokens + kv_cache["k"][:, :num_rolled_tokens] = kv_cache["k"][ :, - sink_tokens + num_evicted_tokens : sink_tokens - + num_evicted_tokens + num_evicted_tokens : num_evicted_tokens + num_rolled_tokens, ].clone() - kv_cache["v"][ - :, sink_tokens : sink_tokens + num_rolled_tokens - ] = kv_cache["v"][ + kv_cache["v"][:, :num_rolled_tokens] = kv_cache["v"][ :, - sink_tokens + num_evicted_tokens : sink_tokens - + num_evicted_tokens + num_evicted_tokens : num_evicted_tokens + num_rolled_tokens, ].clone() local_end_index = ( @@ -360,7 +349,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 +372,6 @@ class CausalMatrixGameTransformerBlock(nn.Module): dim, num_heads, local_attn_size=local_attn_size, - sink_size=sink_size, qk_norm=qk_norm, eps=eps, ) @@ -630,11 +617,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 +649,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, From 404c04934a455e1a41395d9fbf9d5e41164cb35b Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Fri, 27 Feb 2026 04:51:33 +0000 Subject: [PATCH 2/6] extract kv_cache into single method --- .../models/dits/matrixgame/action_module.py | 107 ++---------------- .../models/dits/matrixgame/causal_model.py | 67 ++--------- fastvideo/models/dits/matrixgame/kv_cache.py | 89 +++++++++++++++ 3 files changed, 108 insertions(+), 155 deletions(-) create mode 100644 fastvideo/models/dits/matrixgame/kv_cache.py diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index 3a7b0ff53..e30279876 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -15,6 +15,7 @@ from fastvideo.layers.rotary_embedding import ( _apply_rotary_emb, ) from fastvideo.platforms import AttentionBackendEnum +from .kv_cache import update_kv_cache_and_get_attended_kv DISABLE_COMPILE = False @@ -105,112 +106,26 @@ def _update_kv_cache_and_attend( store_first_only: bool = False, repeat_factor: int | None = None, ) -> 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 - - max_attention_size = local_attn_size - 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"]) + cached_k, cached_v = update_kv_cache_and_get_attended_kv( + kv_cache=kv_cache, + new_k=k, + new_v=v, + current_start=start_frame, + num_new_tokens=num_new_tokens, + max_attn_size=local_attn_size, + store_first_only=store_first_only, ) - 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 - # Roll k cache - kv_cache["k"][:, :num_rolled_tokens] = kv_cache["k"][ - :, num_evicted_tokens : num_evicted_tokens + num_rolled_tokens - ].clone() - # Roll v cache - kv_cache["v"][:, :num_rolled_tokens] = kv_cache["v"][ - :, num_evicted_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 + return attn_layer(q, cached_k, cached_v) class ActionModule(nn.Module): diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 0627b96a6..1d4446aa4 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -44,6 +44,7 @@ from fastvideo.models.dits.wanvideo import ( from fastvideo.platforms import AttentionBackendEnum, current_platform from .action_module import ActionModule +from .kv_cache import update_kv_cache_and_get_attended_kv from .model import MatrixGameCrossAttention logger = init_logger(__name__) @@ -266,63 +267,20 @@ class CausalMatrixGameSelfAttention(nn.Module): "grid_sizes not provided, using q.shape[1] as frame_seqlen" ) - current_end = current_start + roped_query.shape[1] - # 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"]) + k_for_attn, v_for_attn = update_kv_cache_and_get_attended_kv( + kv_cache=kv_cache, + new_k=roped_key, + new_v=v, + current_start=current_start, + num_new_tokens=roped_query.shape[1], + max_attn_size=max_attention_size, ) - 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 - kv_cache["k"][:, :num_rolled_tokens] = kv_cache["k"][ - :, - num_evicted_tokens : num_evicted_tokens - + num_rolled_tokens, - ].clone() - kv_cache["v"][:, :num_rolled_tokens] = kv_cache["v"][ - :, - num_evicted_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), @@ -330,15 +288,6 @@ class CausalMatrixGameSelfAttention(nn.Module): 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 diff --git a/fastvideo/models/dits/matrixgame/kv_cache.py b/fastvideo/models/dits/matrixgame/kv_cache.py new file mode 100644 index 000000000..7962a4fed --- /dev/null +++ b/fastvideo/models/dits/matrixgame/kv_cache.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch + + +def update_kv_cache_and_get_attended_kv( + kv_cache: dict[str, torch.Tensor | int], + new_k: torch.Tensor, + new_v: torch.Tensor, + current_start: int, + num_new_tokens: int, + max_attn_size: int, + sink_tokens: int = 0, + store_first_only: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + current_end = current_start + num_new_tokens + kv_cache_size = kv_cache["k"].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"]) + ) + 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"]) + ) + + 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 + ) + 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 = ( + original_local_end_index + + current_end + - original_global_end_index + - num_evicted_tokens + ) + else: + local_end_index = ( + original_local_end_index + + current_end + - original_global_end_index + ) + local_start_index = local_end_index - num_new_tokens + + if store_first_only: + kv_cache["k"][:, local_start_index:local_end_index] = new_k[:1] + kv_cache["v"][:, local_start_index:local_end_index] = new_v[:1] + else: + kv_cache["k"][:, local_start_index:local_end_index] = new_k + kv_cache["v"][:, local_start_index:local_end_index] = new_v + + cache_start = max(0, local_end_index - max_attn_size) + cached_k = kv_cache["k"][:, cache_start:local_end_index] + cached_v = kv_cache["v"][:, cache_start:local_end_index] + + 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 cached_k, cached_v From a3e6ec1fcac8bc7c3043573a248c8c5fe6df67cb Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Fri, 27 Feb 2026 07:14:29 +0000 Subject: [PATCH 3/6] remove global/local end index --- .../models/dits/matrixgame/action_module.py | 41 +++- .../models/dits/matrixgame/causal_model.py | 36 +++- fastvideo/models/dits/matrixgame/kv_cache.py | 182 +++++++++++------- .../pipelines/stages/matrixgame_denoising.py | 15 +- 4 files changed, 181 insertions(+), 93 deletions(-) diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index e30279876..385b2a61a 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -15,7 +15,10 @@ from fastvideo.layers.rotary_embedding import ( _apply_rotary_emb, ) from fastvideo.platforms import AttentionBackendEnum -from .kv_cache import update_kv_cache_and_get_attended_kv +from .kv_cache import ( + get_attended_kv_without_update, + update_kv_cache_and_get_attended_kv, +) DISABLE_COMPILE = False @@ -105,21 +108,31 @@ 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.""" assert ( 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] - cached_k, cached_v = update_kv_cache_and_get_attended_kv( - kv_cache=kv_cache, - new_k=k, - new_v=v, - current_start=start_frame, - num_new_tokens=num_new_tokens, - max_attn_size=local_attn_size, - store_first_only=store_first_only, - ) + 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) @@ -363,6 +376,7 @@ class ActionModule(nn.Module): 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, @@ -458,6 +472,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) @@ -479,6 +494,7 @@ class ActionModule(nn.Module): 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, @@ -584,6 +600,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) @@ -612,6 +629,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) @@ -635,6 +653,7 @@ class ActionModule(nn.Module): 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 @@ -704,6 +723,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, @@ -727,6 +747,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 1d4446aa4..4f4e7f751 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -44,7 +44,10 @@ from fastvideo.models.dits.wanvideo import ( from fastvideo.platforms import AttentionBackendEnum, current_platform from .action_module import ActionModule -from .kv_cache import update_kv_cache_and_get_attended_kv +from .kv_cache import ( + get_attended_kv_without_update, + update_kv_cache_and_get_attended_kv, +) from .model import MatrixGameCrossAttention logger = init_logger(__name__) @@ -181,6 +184,7 @@ class CausalMatrixGameSelfAttention(nn.Module): kv_cache: dict | 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 @@ -273,14 +277,22 @@ class CausalMatrixGameSelfAttention(nn.Module): if self.local_attn_size == -1 else self.local_attn_size * frame_seqlen ) - k_for_attn, v_for_attn = update_kv_cache_and_get_attended_kv( - kv_cache=kv_cache, - new_k=roped_key, - new_v=v, - current_start=current_start, - num_new_tokens=roped_query.shape[1], - max_attn_size=max_attention_size, - ) + 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), @@ -417,6 +429,7 @@ class CausalMatrixGameTransformerBlock(nn.Module): 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) @@ -464,6 +477,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) @@ -503,6 +517,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( @@ -819,6 +834,7 @@ class CausalMatrixGameWanModel(BaseDiT): 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 @@ -1019,6 +1035,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( @@ -1039,6 +1056,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 index 7962a4fed..5008d2eb8 100644 --- a/fastvideo/models/dits/matrixgame/kv_cache.py +++ b/fastvideo/models/dits/matrixgame/kv_cache.py @@ -3,87 +3,139 @@ import torch +def _store_first_tokens( + new_tensor: torch.Tensor, + num_new_tokens: int, + target_batch: int, +) -> torch.Tensor: + first = new_tensor[:1] + first = first.expand(target_batch, num_new_tokens, -1, -1) + return first + + +def get_attended_kv_without_update( + kv_cache: dict[str, torch.Tensor | int], + new_k: torch.Tensor, + new_v: torch.Tensor, + num_new_tokens: int, + max_attn_size: int, + sink_tokens: int = 0, + 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"]) + ) + + if store_first_only: + write_k = _store_first_tokens(new_k, num_new_tokens, target_batch) + write_v = _store_first_tokens(new_v, num_new_tokens, target_batch) + else: + write_k = new_k + write_v = new_v + + 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 + if sink_tokens > 0: + kept_k = torch.cat( + [ + cache_k[:, :sink_tokens], + cache_k[:, sink_tokens + num_evicted_tokens :], + ], + dim=1, + ) + kept_v = torch.cat( + [ + cache_v[:, :sink_tokens], + cache_v[:, sink_tokens + num_evicted_tokens :], + ], + dim=1, + ) + else: + kept_k = cache_k[:, num_evicted_tokens:] + kept_v = cache_v[:, num_evicted_tokens:] + attend_k = torch.cat([kept_k, write_k], dim=1) + attend_v = torch.cat([kept_v, write_v], dim=1) + else: + attend_k = torch.cat([cache_k, write_k], dim=1) + attend_v = torch.cat([cache_v, write_v], dim=1) + + attend_length = attend_k.shape[1] + attn_start = max(0, attend_length - max_attn_size) + cached_k = attend_k[:, attn_start:attend_length] + cached_v = attend_v[:, attn_start:attend_length] + return cached_k, cached_v + + def update_kv_cache_and_get_attended_kv( kv_cache: dict[str, torch.Tensor | int], new_k: torch.Tensor, new_v: torch.Tensor, - current_start: int, num_new_tokens: int, max_attn_size: int, sink_tokens: int = 0, store_first_only: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - current_end = current_start + num_new_tokens - kv_cache_size = kv_cache["k"].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"]) - ) - 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"]) + 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"]) ) - 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 - ) - 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 = ( - original_local_end_index - + current_end - - original_global_end_index - - num_evicted_tokens - ) - else: - local_end_index = ( - original_local_end_index - + current_end - - original_global_end_index - ) - local_start_index = local_end_index - num_new_tokens + cache_k = kv_cache["k"][:, :cache_length] + cache_v = kv_cache["v"][:, :cache_length] if store_first_only: - kv_cache["k"][:, local_start_index:local_end_index] = new_k[:1] - kv_cache["v"][:, local_start_index:local_end_index] = new_v[:1] + 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]) else: - kv_cache["k"][:, local_start_index:local_end_index] = new_k - kv_cache["v"][:, local_start_index:local_end_index] = new_v + write_k = new_k + write_v = new_v - cache_start = max(0, local_end_index - max_attn_size) - cached_k = kv_cache["k"][:, cache_start:local_end_index] - cached_v = kv_cache["v"][:, cache_start:local_end_index] + if num_new_tokens + cache_length > kv_capacity: + num_evicted_tokens = num_new_tokens + cache_length - kv_capacity + if sink_tokens > 0: + kept_k = torch.cat( + [ + cache_k[:, :sink_tokens], + cache_k[:, sink_tokens + num_evicted_tokens :], + ], + dim=1, + ) + kept_v = torch.cat( + [ + cache_v[:, :sink_tokens], + cache_v[:, sink_tokens + num_evicted_tokens :], + ], + dim=1, + ) + else: + kept_k = cache_k[:, num_evicted_tokens:] + kept_v = cache_v[:, num_evicted_tokens:] + updated_k = torch.cat([kept_k, write_k], dim=1) + updated_v = torch.cat([kept_v, write_v], dim=1) + else: + updated_k = torch.cat([cache_k, write_k], dim=1) + updated_v = torch.cat([cache_v, write_v], dim=1) - if isinstance(kv_cache["global_end_index"], torch.Tensor): - kv_cache["global_end_index"].fill_(current_end) + 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] + + 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] + + if isinstance(kv_cache["length"], torch.Tensor): + kv_cache["length"].fill_(updated_length) 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 + kv_cache["length"] = updated_length return cached_k, cached_v diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index cd9599bbb..39d74d19c 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -323,9 +323,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): ], dtype=dtype, device=device), - "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), - "local_end_index": + "length": torch.tensor([0], dtype=torch.long, device=device), }) @@ -361,9 +359,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): ], dtype=dtype, device=device), - "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), - "local_end_index": + "length": torch.tensor([0], dtype=torch.long, device=device), }) kv_cache_mouse.append({ @@ -381,9 +377,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): ], dtype=dtype, device=device), - "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), - "local_end_index": + "length": torch.tensor([0], dtype=torch.long, device=device), }) @@ -495,6 +489,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 +600,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 +621,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, ) From 1d5b6962a542aaeff541fe32cfa05f5cf398c332 Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Fri, 27 Feb 2026 07:46:24 +0000 Subject: [PATCH 4/6] remove sinktoken --- fastvideo/models/dits/matrixgame/kv_cache.py | 42 ++------------------ 1 file changed, 4 insertions(+), 38 deletions(-) diff --git a/fastvideo/models/dits/matrixgame/kv_cache.py b/fastvideo/models/dits/matrixgame/kv_cache.py index 5008d2eb8..8a78d88da 100644 --- a/fastvideo/models/dits/matrixgame/kv_cache.py +++ b/fastvideo/models/dits/matrixgame/kv_cache.py @@ -19,7 +19,6 @@ def get_attended_kv_without_update( new_v: torch.Tensor, num_new_tokens: int, max_attn_size: int, - sink_tokens: int = 0, store_first_only: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: kv_capacity = kv_cache["k"].shape[1] @@ -42,24 +41,8 @@ def get_attended_kv_without_update( if num_new_tokens + cache_length > kv_capacity: num_evicted_tokens = num_new_tokens + cache_length - kv_capacity - if sink_tokens > 0: - kept_k = torch.cat( - [ - cache_k[:, :sink_tokens], - cache_k[:, sink_tokens + num_evicted_tokens :], - ], - dim=1, - ) - kept_v = torch.cat( - [ - cache_v[:, :sink_tokens], - cache_v[:, sink_tokens + num_evicted_tokens :], - ], - dim=1, - ) - else: - kept_k = cache_k[:, num_evicted_tokens:] - kept_v = cache_v[:, num_evicted_tokens:] + kept_k = cache_k[:, num_evicted_tokens:] + kept_v = cache_v[:, num_evicted_tokens:] attend_k = torch.cat([kept_k, write_k], dim=1) attend_v = torch.cat([kept_v, write_v], dim=1) else: @@ -79,7 +62,6 @@ def update_kv_cache_and_get_attended_kv( new_v: torch.Tensor, num_new_tokens: int, max_attn_size: int, - sink_tokens: int = 0, store_first_only: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: kv_capacity = kv_cache["k"].shape[1] @@ -101,24 +83,8 @@ def update_kv_cache_and_get_attended_kv( if num_new_tokens + cache_length > kv_capacity: num_evicted_tokens = num_new_tokens + cache_length - kv_capacity - if sink_tokens > 0: - kept_k = torch.cat( - [ - cache_k[:, :sink_tokens], - cache_k[:, sink_tokens + num_evicted_tokens :], - ], - dim=1, - ) - kept_v = torch.cat( - [ - cache_v[:, :sink_tokens], - cache_v[:, sink_tokens + num_evicted_tokens :], - ], - dim=1, - ) - else: - kept_k = cache_k[:, num_evicted_tokens:] - kept_v = cache_v[:, num_evicted_tokens:] + kept_k = cache_k[:, num_evicted_tokens:] + kept_v = cache_v[:, num_evicted_tokens:] updated_k = torch.cat([kept_k, write_k], dim=1) updated_v = torch.cat([kept_v, write_v], dim=1) else: From 8cae20a654f4057fabe5010e690ee75d8450fedd Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Fri, 27 Feb 2026 08:29:02 +0000 Subject: [PATCH 5/6] kv cache dataclass --- .../models/dits/matrixgame/action_module.py | 50 ++---- .../models/dits/matrixgame/causal_model.py | 56 +++--- fastvideo/models/dits/matrixgame/kv_cache.py | 162 ++++++++++++++---- .../pipelines/stages/matrixgame_denoising.py | 93 ++++------ 4 files changed, 211 insertions(+), 150 deletions(-) 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 From 40b681c05b6b04a591e6c7925d34c11723c8ad5a Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Fri, 27 Feb 2026 00:54:47 -0800 Subject: [PATCH 6/6] unified everything --- fastvideo/models/dits/matrixgame/kv_cache.py | 185 +++++++------------ 1 file changed, 69 insertions(+), 116 deletions(-) diff --git a/fastvideo/models/dits/matrixgame/kv_cache.py b/fastvideo/models/dits/matrixgame/kv_cache.py index be05c5af6..6d9f9f67b 100644 --- a/fastvideo/models/dits/matrixgame/kv_cache.py +++ b/fastvideo/models/dits/matrixgame/kv_cache.py @@ -55,11 +55,73 @@ class KVCache: def batch_size(self) -> int: return self.k.shape[0] - def length_int(self) -> int: + def get_len(self) -> int: return int(self.length.item()) - def set_length(self, value: int) -> None: - self.length.fill_(int(value)) + 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 @@ -76,99 +138,6 @@ class KVCacheDict: } -def _store_first_tokens( - new_tensor: torch.Tensor, - num_new_tokens: int, - target_batch: int, -) -> torch.Tensor: - first = new_tensor[:1] - first = first.expand(target_batch, num_new_tokens, -1, -1) - return first - - -def get_attended_kv_without_update( - 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.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) - write_v = _store_first_tokens(new_v, num_new_tokens, target_batch) - else: - write_k = new_k - write_v = new_v - - 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 - kept_k = cache_k[:, num_evicted_tokens:] - kept_v = cache_v[:, num_evicted_tokens:] - attend_k = torch.cat([kept_k, write_k], dim=1) - attend_v = torch.cat([kept_v, write_v], dim=1) - else: - attend_k = torch.cat([cache_k, write_k], dim=1) - attend_v = torch.cat([cache_v, write_v], dim=1) - - attend_length = attend_k.shape[1] - attn_start = max(0, attend_length - max_attn_size) - cached_k = attend_k[:, attn_start:attend_length] - cached_v = attend_v[:, attn_start:attend_length] - return cached_k, cached_v - - -def update_kv_cache_and_get_attended_kv( - 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.capacity - cache_length = kv_cache.length_int() - - 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.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 - - if num_new_tokens + cache_length > kv_capacity: - num_evicted_tokens = num_new_tokens + cache_length - kv_capacity - kept_k = cache_k[:, num_evicted_tokens:] - kept_v = cache_v[:, num_evicted_tokens:] - updated_k = torch.cat([kept_k, write_k], dim=1) - updated_v = torch.cat([kept_v, write_v], dim=1) - else: - updated_k = torch.cat([cache_k, write_k], dim=1) - 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] - - 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] - - kv_cache.set_length(updated_length) - - return cached_k, cached_v - - def attend_with_kv_cache( q: torch.Tensor, new_k: torch.Tensor, @@ -183,29 +152,13 @@ def attend_with_kv_cache( 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, + 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 = 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, + cached_k, cached_v = kv_cache.get_view( + new_k, new_v, max_attn_size, store_first_only, repeat_factor, num_new_tokens ) - - 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)