remove global/local end index

2026-02-27 07:14:29 +00:00
parent 404c04934a
commit a3e6ec1fca
4 changed files with 181 additions and 93 deletions
@@ -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,
@@ -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(
@@ -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
@@ -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,
                )