nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

kv cache dataclass

Browse Source
This commit is contained in:
H1yori233
2026-02-27 08:29:02 +00:00
parent 1d5b6962a5
commit 8cae20a654
4 changed files with 211 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/dits/matrixgame/action_module.py
+19 -31
View File
@@ -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,
fastvideo/models/dits/matrixgame/causal_model.py
+24 -32
View File
@@ -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,
fastvideo/models/dits/matrixgame/kv_cache.py
+133 -29
View File
@@ -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)
fastvideo/pipelines/stages/matrixgame_denoising.py
+35 -58
View File
@@ -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