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

[perf] causal MatrixGame optimization (#1078)

Browse Source
This commit is contained in:
XOR-op
2026-02-15 09:51:00 +08:00
committed by GitHub
parent 36bf37e9ba
commit 959dbcc8a2
6 changed files with 64 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/dits/causal_wanvideo.py
+9 -9
View File
@@ -126,26 +126,26 @@ class CausalWanSelfAttention(nn.Module):
# If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache
kv_cache_size = kv_cache["k"].shape[1]
num_new_tokens = roped_query.shape[1]
if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"]) and (
num_new_tokens + kv_cache["local_end_index"] > kv_cache_size):
# Calculate the number of new tokens added in this step
# Shift existing cache content left to discard oldest tokens
# Clone the source slice to avoid overlapping memory error
num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"] - kv_cache_size
num_rolled_tokens = kv_cache["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()
# Insert the new keys/values at the end
local_end_index = kv_cache["local_end_index"].item() + current_end - \
kv_cache["global_end_index"].item() - num_evicted_tokens
local_end_index = kv_cache["local_end_index"] + current_end - \
kv_cache["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:
# Assign new keys/values directly up to current_end
local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
local_end_index = kv_cache["local_end_index"] + current_end - kv_cache["global_end_index"]
local_start_index = local_end_index - num_new_tokens
kv_cache["k"] = kv_cache["k"].detach()
kv_cache["v"] = kv_cache["v"].detach()
@@ -157,8 +157,8 @@ class CausalWanSelfAttention(nn.Module):
kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index],
kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index]
)
kv_cache["global_end_index"].fill_(current_end)
kv_cache["local_end_index"].fill_(local_end_index)
kv_cache["global_end_index"] = current_end
kv_cache["local_end_index"] = local_end_index
return x
fastvideo/models/dits/matrixgame/action_module.py
+14 -11
View File
@@ -95,7 +95,7 @@ def _update_kv_cache_and_attend(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
kv_cache: dict[str, torch.Tensor],
kv_cache: dict[str, torch.Tensor | int],
attn_layer: LocalAttention,
start_frame: int,
num_frame_per_block: int,
@@ -139,15 +139,18 @@ def _update_kv_cache_and_attend(
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 = kv_cache["global_end_index"]
original_local_end_index = kv_cache["local_end_index"]
# Check if we need to evict tokens
if (current_end > kv_cache["global_end_index"].item()) and (
num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size
if (current_end > original_global_end_index) and (
num_new_tokens + original_local_end_index > kv_cache_size
):
num_evicted_tokens = (
num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
num_new_tokens + original_local_end_index - kv_cache_size
)
num_rolled_tokens = (
kv_cache["local_end_index"].item()
original_local_end_index
- num_evicted_tokens
- sink_tokens
)
@@ -171,18 +174,18 @@ def _update_kv_cache_and_attend(
)
# Calculate indices with eviction adjustment
local_end_index = (
kv_cache["local_end_index"].item()
original_local_end_index
+ current_end
- kv_cache["global_end_index"].item()
- original_global_end_index
- num_evicted_tokens
)
local_start_index = local_end_index - num_new_tokens
else:
# Calculate indices without eviction
local_end_index = (
kv_cache["local_end_index"].item()
original_local_end_index
+ current_end
- kv_cache["global_end_index"].item()
- original_global_end_index
)
local_start_index = local_end_index - num_new_tokens
@@ -206,8 +209,8 @@ def _update_kv_cache_and_attend(
attn = attn_layer(q, cached_k, cached_v)
# Update indices
kv_cache["global_end_index"].fill_(current_end)
kv_cache["local_end_index"].fill_(local_end_index)
kv_cache["global_end_index"] = current_end
kv_cache["local_end_index"] = local_end_index
return attn
fastvideo/models/dits/matrixgame/causal_model.py
+28 -35
View File
@@ -49,7 +49,9 @@ from .model import MatrixGameCrossAttention
logger = init_logger(__name__)
def causal_rope_apply(x, grid_sizes, freqs, start_frame=0):
def causal_rope_apply(
x: torch.Tensor, grid_sizes: tuple[int, int, int], freqs: torch.Tensor, start_frame: int = 0
):
n, c = x.size(2), x.size(3) // 2
# split freqs
@@ -57,7 +59,7 @@ def causal_rope_apply(x, grid_sizes, freqs, start_frame=0):
# loop over samples
output = []
f, h, w = grid_sizes.tolist()
f, h, w = grid_sizes
for i in range(len(x)):
seq_len = f * h * w
@@ -176,16 +178,16 @@ class CausalMatrixGameSelfAttention(nn.Module):
v: torch.Tensor,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
block_mask: BlockMask,
grid_sizes: tuple[int, int, int],
kv_cache: dict | None = None,
current_start: int = 0,
cache_start: int | None = None,
grid_sizes: torch.Tensor | None = None,
):
if cache_start is None:
cache_start = current_start
# Calculate start_frame for causal mode
if kv_cache is not None and grid_sizes is not None:
if kv_cache is not None:
frame_seqlen = int(grid_sizes[1] * grid_sizes[2])
start_frame = current_start // frame_seqlen
else:
@@ -279,17 +281,14 @@ class CausalMatrixGameSelfAttention(nn.Module):
kv_cache_size = kv_cache["k"].shape[1]
num_new_tokens = roped_query.shape[1]
if (current_end > kv_cache["global_end_index"].item()) and (
num_new_tokens + kv_cache["local_end_index"].item()
> kv_cache_size
if (current_end > kv_cache["global_end_index"]) and (
num_new_tokens + kv_cache["local_end_index"] > kv_cache_size
):
num_evicted_tokens = (
num_new_tokens
+ kv_cache["local_end_index"].item()
- kv_cache_size
num_new_tokens + kv_cache["local_end_index"] - kv_cache_size
)
num_rolled_tokens = (
kv_cache["local_end_index"].item()
kv_cache["local_end_index"]
- num_evicted_tokens
- sink_tokens
)
@@ -310,9 +309,9 @@ class CausalMatrixGameSelfAttention(nn.Module):
+ num_rolled_tokens,
].clone()
local_end_index = (
kv_cache["local_end_index"].item()
kv_cache["local_end_index"]
+ current_end
- kv_cache["global_end_index"].item()
- kv_cache["global_end_index"]
- num_evicted_tokens
)
local_start_index = local_end_index - num_new_tokens
@@ -320,9 +319,9 @@ class CausalMatrixGameSelfAttention(nn.Module):
kv_cache["v"][:, local_start_index:local_end_index] = v
else:
local_end_index = (
kv_cache["local_end_index"].item()
kv_cache["local_end_index"]
+ current_end
- kv_cache["global_end_index"].item()
- kv_cache["global_end_index"]
)
local_start_index = local_end_index - num_new_tokens
@@ -339,8 +338,8 @@ class CausalMatrixGameSelfAttention(nn.Module):
v_for_attn.transpose(1, 2),
).transpose(1, 2)
kv_cache["global_end_index"].fill_(current_end)
kv_cache["local_end_index"].fill_(local_end_index)
kv_cache["global_end_index"] = current_end
kv_cache["local_end_index"] = local_end_index
return x
@@ -460,7 +459,7 @@ class CausalMatrixGameTransformerBlock(nn.Module):
temb: torch.Tensor,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
block_mask: BlockMask,
grid_sizes: torch.Tensor,
grid_sizes: tuple[int, int, int],
mouse_cond: torch.Tensor | None = None,
keyboard_cond: torch.Tensor | None = None,
block_mask_mouse: BlockMask | None = None,
@@ -516,10 +515,10 @@ class CausalMatrixGameTransformerBlock(nn.Module):
value,
freqs_cis,
block_mask,
grid_sizes,
kv_cache,
current_start,
cache_start,
grid_sizes,
)
attn_output = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output)
@@ -546,9 +545,9 @@ class CausalMatrixGameTransformerBlock(nn.Module):
hidden_states = self.action_model(
hidden_states,
int(grid_sizes[0]),
int(grid_sizes[1]),
int(grid_sizes[2]),
grid_sizes[0],
grid_sizes[1],
grid_sizes[2],
mouse_cond,
keyboard_cond,
block_mask_mouse,
@@ -981,10 +980,7 @@ class CausalMatrixGameWanModel(BaseDiT):
freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
hidden_states = self.patch_embedding(hidden_states)
grid_sizes = torch.tensor(
[post_patch_num_frames, post_patch_height, post_patch_width],
device=hidden_states.device,
)
grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)
hidden_states = hidden_states.flatten(2).transpose(1, 2)
if timestep.dim() == 2:
@@ -1200,9 +1196,9 @@ class CausalMatrixGameWanModel(BaseDiT):
hidden_states.shape
)
p_t, p_h, p_w = self.patch_size
post_patch_num_frames = num_frames // p_t
post_patch_height = height // p_h
post_patch_width = width // p_w
post_patch_num_frames: int = num_frames // p_t
post_patch_height: int = height // p_h
post_patch_width: int = width // p_w
d = self.hidden_size // self.num_attention_heads
rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
@@ -1263,10 +1259,7 @@ class CausalMatrixGameWanModel(BaseDiT):
)
hidden_states = self.patch_embedding(hidden_states)
grid_sizes = torch.tensor(
[post_patch_num_frames, post_patch_height, post_patch_width],
device=hidden_states.device,
)
grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width)
hidden_states = hidden_states.flatten(2).transpose(1, 2)
if timestep.dim() == 2:
@@ -1393,10 +1386,10 @@ class CausalMatrixGameWanModel(BaseDiT):
else:
return self._forward_train(*args, **kwargs)
def unpatchify(self, x, grid_sizes):
def unpatchify(self, x: torch.Tensor, grid_sizes: tuple[int, int, int]) -> torch.Tensor:
c = self.proj_out.out_features // math.prod(self.patch_size)
p_t, p_h, p_w = self.patch_size
f, h, w = grid_sizes.tolist()
f, h, w = grid_sizes
x = x[:, : f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c)
x = x.permute(0, 7, 1, 4, 2, 5, 3, 6)
fastvideo/pipelines/stages/causal_denoising.py
+3 -3
View File
@@ -436,9 +436,9 @@ class CausalDMDDenosingStage(DenoisingStage):
dtype=dtype,
device=device),
"global_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
"local_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
})
return kv_cache1
@@ -494,4 +494,4 @@ class CausalDMDDenosingStage(DenoisingStage):
result.add_check(
"negative_prompt_embeds", batch.negative_prompt_embeds, lambda x:
not batch.do_classifier_free_guidance or V.list_not_empty(x))
return result
return result
fastvideo/pipelines/stages/matrixgame_denoising.py
+6 -6
View File
@@ -335,9 +335,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
dtype=dtype,
device=device),
"global_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
"local_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
})
return kv_cache
@@ -373,9 +373,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
dtype=dtype,
device=device),
"global_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
"local_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
})
kv_cache_mouse.append({
"k":
@@ -393,9 +393,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage):
dtype=dtype,
device=device),
"global_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
"local_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
})
return kv_cache_mouse, kv_cache_keyboard
fastvideo/training/self_forcing_distillation_pipeline.py
+4 -4
View File
@@ -582,9 +582,9 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
dtype=dtype,
device=device),
"global_end_index":
torch.tensor([0], dtype=torch.long, device=device),
0,
"local_end_index":
torch.tensor([0], dtype=torch.long, device=device)
0
})
# Initialize cross-attention cache
@@ -616,8 +616,8 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
"""Reset KV cache and cross-attention cache to clean state."""
if kv_cache is not None:
for cache_dict in kv_cache:
cache_dict["global_end_index"].fill_(0)
cache_dict["local_end_index"].fill_(0)
cache_dict["global_end_index"] = 0
cache_dict["local_end_index"] = 0
cache_dict["k"].zero_()
cache_dict["v"].zero_()