From a3e6ec1fcac8bc7c3043573a248c8c5fe6df67cb Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Fri, 27 Feb 2026 07:14:29 +0000 Subject: [PATCH] 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, )