cleanup and optimize
This commit is contained in:
kijai
@@ -15,6 +15,10 @@ try:
|
||||
except:
|
||||
BlockMask = create_block_mask = flex_attention = None
|
||||
pass
|
||||
try:
|
||||
from ..radial_attention.attn_mask import RadialAttention
|
||||
except:
|
||||
pass
|
||||
|
||||
from .attention import attention
|
||||
import numpy as np
|
||||
@@ -262,7 +266,7 @@ class WanSelfAttention(nn.Module):
|
||||
|
||||
#radial attention
|
||||
self.layer_idx = layer_idx
|
||||
self.dense_timestep = 12
|
||||
self.dense_timestep = 10
|
||||
self.dense_block = 1
|
||||
self.decay_factor = 0.2
|
||||
self.sparse_type = "radial"
|
||||
@@ -317,22 +321,11 @@ class WanSelfAttention(nn.Module):
|
||||
block_mask=block_mask
|
||||
)[:, :, :-padded_length].transpose(2, 1)
|
||||
elif self.attention_mode == 'radial_sage_attention':
|
||||
from ..radial_attention.attn_mask import RadialAttention
|
||||
batch_size = q.shape[0]
|
||||
q = rearrange(q, "b s h d -> (b s) h d")
|
||||
k = rearrange(k, "b s h d -> (b s) h d")
|
||||
v = rearrange(v, "b s h d -> (b s) h d")
|
||||
if current_step < self.dense_timestep or self.layer_idx < self.dense_block or self.sparse_type == "dense":
|
||||
x = RadialAttention(
|
||||
query=q, key=k, value=v, mask_map=self.mask_map, sparsity_type="dense", block_size=128, decay_factor=self.decay_factor, model_type="wan", pre_defined_mask=None
|
||||
)
|
||||
dense_step = current_step < self.dense_timestep or self.layer_idx < self.dense_block or self.sparse_type == "dense"
|
||||
if dense_step:
|
||||
x = RadialAttention(query=q, key=k, value=v, mask_map=self.mask_map, sparsity_type="dense", block_size=128, decay_factor=self.decay_factor)
|
||||
else:
|
||||
# apply radial attention
|
||||
x = RadialAttention(
|
||||
query=q, key=k, value=v, mask_map=self.mask_map, sparsity_type="radial", block_size=128, decay_factor=self.decay_factor, model_type="wan", pre_defined_mask=None
|
||||
)
|
||||
x = rearrange(x, "(b s) h d -> b s h d", b=batch_size)
|
||||
|
||||
x = RadialAttention(query=q, key=k, value=v, mask_map=self.mask_map, sparsity_type="radial", block_size=128, decay_factor=self.decay_factor)
|
||||
else:
|
||||
x = attention(
|
||||
q, k, v,
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
import torch
|
||||
#import flashinfer
|
||||
import matplotlib.pyplot as plt
|
||||
try:
|
||||
from sparse_sageattn import sparse_sageattn
|
||||
except:
|
||||
sparse_sageattn = None
|
||||
raise ImportError("Package is not installed: https://github.com/jt-zhang/Sparse_SageAttention_API")
|
||||
from einops import rearrange, repeat
|
||||
|
||||
def sparge_mask_convert(mask: torch.Tensor, block_size: int = 128) -> torch.Tensor:
|
||||
@@ -21,25 +20,28 @@ def sparge_mask_convert(mask: torch.Tensor, block_size: int = 128) -> torch.Tens
|
||||
|
||||
return new_mask
|
||||
|
||||
def get_indptr_from_mask(mask, query):
|
||||
from comfy import model_management as mm
|
||||
device = mm.get_torch_device()
|
||||
|
||||
def get_indptr_from_mask(mask):
|
||||
# query shows the device of the indptr
|
||||
# indptr (torch.Tensor) - the block index pointer of the block-sparse matrix on row dimension,
|
||||
# shape `(MB + 1,)`, where `MB` is the number of blocks in the row dimension.
|
||||
# The first element is always 0, and the last element is the number of blocks in the row dimension.
|
||||
# The rest of the elements are the number of blocks in each row.
|
||||
# the mask is already a block sparse mask
|
||||
indptr = torch.zeros(mask.shape[0] + 1, device=query.device, dtype=torch.int32)
|
||||
indptr = torch.zeros(mask.shape[0] + 1, device=device, dtype=torch.int32)
|
||||
indptr[0] = 0
|
||||
row_counts = mask.sum(dim=1).flatten() # Ensure 1D output [num_blocks_row]
|
||||
indptr[1:] = torch.cumsum(row_counts, dim=0)
|
||||
return indptr
|
||||
|
||||
def get_indices_from_mask(mask, query):
|
||||
def get_indices_from_mask(mask):
|
||||
# indices (torch.Tensor) - the block indices of the block-sparse matrix on column dimension,
|
||||
# shape `(nnz,),` where `nnz` is the number of non-zero blocks.
|
||||
# The elements in `indices` array should be less than `NB`: the number of blocks in the column dimension.
|
||||
nonzero_indices = torch.nonzero(mask)
|
||||
indices = nonzero_indices[:, 1].to(dtype=torch.int32, device=query.device)
|
||||
indices = nonzero_indices[:, 1].to(dtype=torch.int32, device=device)
|
||||
return indices
|
||||
|
||||
def shrinkMaskStrict(mask, block_size=128):
|
||||
@@ -71,35 +73,29 @@ def pad_qkv(input_tensor, block_size=128):
|
||||
|
||||
return padded_tensor
|
||||
|
||||
def get_diagonal_split_mask(i, j, token_per_frame, sparse_type, query):
|
||||
def get_diagonal_split_mask(i, j, token_per_frame, sparse_type):
|
||||
assert(sparse_type in ["radial"])
|
||||
dist = abs(i - j)
|
||||
group = dist.bit_length()
|
||||
threshold = 128 # hardcoded threshold for now, which is equal to block-size
|
||||
decay_length = 2 ** token_per_frame.bit_length() / 2 ** group
|
||||
if decay_length >= threshold:
|
||||
return torch.ones((token_per_frame, token_per_frame), device=query.device, dtype=torch.bool)
|
||||
return torch.ones((token_per_frame, token_per_frame), device=device, dtype=torch.bool)
|
||||
|
||||
split_factor = int(threshold / decay_length)
|
||||
modular = dist % split_factor
|
||||
if modular == 0:
|
||||
return torch.ones((token_per_frame, token_per_frame), device=query.device, dtype=torch.bool)
|
||||
return torch.ones((token_per_frame, token_per_frame), device=device, dtype=torch.bool)
|
||||
else:
|
||||
return torch.zeros((token_per_frame, token_per_frame), device=query.device, dtype=torch.bool)
|
||||
return torch.zeros((token_per_frame, token_per_frame), device=device, dtype=torch.bool)
|
||||
|
||||
def get_window_width(i, j, token_per_frame, sparse_type, num_frame, decay_factor=1, block_size=128, model_type=None):
|
||||
def get_window_width(i, j, token_per_frame, sparse_type, decay_factor=1, block_size=128):
|
||||
assert(sparse_type in ["radial"])
|
||||
dist = abs(i - j)
|
||||
if model_type == "wan":
|
||||
if dist < 1:
|
||||
return token_per_frame
|
||||
if dist == 1:
|
||||
return token_per_frame // 2
|
||||
elif model_type == "hunyuan":
|
||||
if dist <= 1:
|
||||
return token_per_frame
|
||||
else:
|
||||
raise ValueError(f"Unknown model type: {model_type}")
|
||||
if dist < 1:
|
||||
return token_per_frame
|
||||
if dist == 1:
|
||||
return token_per_frame // 2
|
||||
group = dist.bit_length()
|
||||
decay_length = 2 ** token_per_frame.bit_length() / 2 ** group * decay_factor
|
||||
threshold = block_size
|
||||
@@ -108,28 +104,28 @@ def get_window_width(i, j, token_per_frame, sparse_type, num_frame, decay_factor
|
||||
else:
|
||||
return threshold
|
||||
|
||||
def gen_log_mask_shrinked(query, s, video_token_num, num_frame, block_size=128, sparse_type="log", decay_factor=0.5, model_type=None):
|
||||
def gen_log_mask_shrinked(device, s, video_token_num, num_frame, block_size=128, sparse_type="log", decay_factor=0.5):
|
||||
"""
|
||||
A more memory friendly version, we generate the attention mask of each frame pair at a time,
|
||||
shrinks it, and stores it into the final result
|
||||
"""
|
||||
final_log_mask = torch.zeros((s // block_size, s // block_size), device=query.device, dtype=torch.bool)
|
||||
final_log_mask = torch.zeros((s // block_size, s // block_size), device=device, dtype=torch.bool)
|
||||
token_per_frame = video_token_num // num_frame
|
||||
video_text_border = video_token_num // block_size
|
||||
|
||||
col_indices = torch.arange(0, token_per_frame, device=query.device).view(1, -1)
|
||||
row_indices = torch.arange(0, token_per_frame, device=query.device).view(-1, 1)
|
||||
col_indices = torch.arange(0, token_per_frame, device=device).view(1, -1)
|
||||
row_indices = torch.arange(0, token_per_frame, device=device).view(-1, 1)
|
||||
final_log_mask[video_text_border:] = True
|
||||
final_log_mask[:, video_text_border:] = True
|
||||
for i in range(num_frame):
|
||||
for j in range(num_frame):
|
||||
local_mask = torch.zeros((token_per_frame, token_per_frame), device=query.device, dtype=torch.bool)
|
||||
if j == 0 and model_type == "wan": # this is attention sink
|
||||
local_mask = torch.ones((token_per_frame, token_per_frame), device=query.device, dtype=torch.bool)
|
||||
local_mask = torch.zeros((token_per_frame, token_per_frame), device=device, dtype=torch.bool)
|
||||
if j == 0: # this is attention sink
|
||||
local_mask = torch.ones((token_per_frame, token_per_frame), device=device, dtype=torch.bool)
|
||||
else:
|
||||
window_width = get_window_width(i, j, token_per_frame, sparse_type, num_frame, decay_factor=decay_factor, block_size=block_size, model_type=model_type)
|
||||
window_width = get_window_width(i, j, token_per_frame, sparse_type, decay_factor=decay_factor, block_size=block_size)
|
||||
local_mask = torch.abs(col_indices - row_indices) <= window_width
|
||||
split_mask = get_diagonal_split_mask(i, j, token_per_frame, sparse_type, query)
|
||||
split_mask = get_diagonal_split_mask(i, j, token_per_frame, sparse_type)
|
||||
local_mask = torch.logical_and(local_mask, split_mask)
|
||||
|
||||
remainder_row = (i * token_per_frame) % block_size
|
||||
@@ -137,7 +133,7 @@ def gen_log_mask_shrinked(query, s, video_token_num, num_frame, block_size=128,
|
||||
# get the padded size
|
||||
all_length_row = remainder_row + ((token_per_frame - 1) // block_size + 1) * block_size
|
||||
all_length_col = remainder_col + ((token_per_frame - 1) // block_size + 1) * block_size
|
||||
padded_local_mask = torch.zeros((all_length_row, all_length_col), device=query.device, dtype=torch.bool)
|
||||
padded_local_mask = torch.zeros((all_length_row, all_length_col), device=device, dtype=torch.bool)
|
||||
padded_local_mask[remainder_row:remainder_row + token_per_frame, remainder_col:remainder_col + token_per_frame] = local_mask
|
||||
# shrink the mask
|
||||
block_mask = shrinkMaskStrict(padded_local_mask, block_size=block_size)
|
||||
@@ -148,7 +144,7 @@ def gen_log_mask_shrinked(query, s, video_token_num, num_frame, block_size=128,
|
||||
block_col_end = block_col_start + block_mask.shape[1]
|
||||
final_log_mask[block_row_start:block_row_end, block_col_start:block_col_end] = torch.logical_or(
|
||||
final_log_mask[block_row_start:block_row_end, block_col_start:block_col_end], block_mask)
|
||||
print(f"mask sparsity: {1 - final_log_mask.sum() / final_log_mask.numel()}")
|
||||
#print(f"mask sparsity: {1 - final_log_mask.sum() / final_log_mask.numel()}")
|
||||
return final_log_mask
|
||||
|
||||
class MaskMap:
|
||||
@@ -158,168 +154,47 @@ class MaskMap:
|
||||
self.num_frame = num_frame
|
||||
self.log_mask = None
|
||||
|
||||
def queryLogMask(self, query, sparse_type, block_size=128, decay_factor=0.5, model_type=None):
|
||||
log_mask = torch.ones((query.shape[0] // block_size, query.shape[0] // block_size), device=query.device, dtype=torch.bool)
|
||||
def queryLogMask(self, seq_len, sparse_type, block_size=128, decay_factor=0.5):
|
||||
log_mask = torch.ones((seq_len // block_size, seq_len // block_size), device=device, dtype=torch.bool)
|
||||
if self.log_mask is None:
|
||||
self.log_mask = gen_log_mask_shrinked(query, query.shape[0], self.video_token_num, self.num_frame, sparse_type=sparse_type, decay_factor=decay_factor, model_type=model_type, block_size=block_size)
|
||||
self.log_mask = gen_log_mask_shrinked(device, seq_len, self.video_token_num, self.num_frame, sparse_type=sparse_type, decay_factor=decay_factor, block_size=block_size)
|
||||
block_bound = self.video_token_num // block_size
|
||||
log_mask[:block_bound, :block_bound] = self.log_mask[:block_bound, :block_bound]
|
||||
return log_mask
|
||||
|
||||
def SpargeSageAttnBackend(query, key, value, mask_map=None, video_mask=None, pre_defined_mask=None, block_size=128):
|
||||
def SpargeSageAttnBackend(query, key, value, mask_map=None, video_mask=None, block_size=128):
|
||||
if video_mask.all():
|
||||
# dense case
|
||||
kv_border = pre_defined_mask[0].sum() if pre_defined_mask is not None else key.shape[0]
|
||||
output_video = sparse_sageattn(
|
||||
query[:mask_map.video_token_num].unsqueeze(0),
|
||||
key[:kv_border, :, :].unsqueeze(0),
|
||||
value[:kv_border, :, :].unsqueeze(0),
|
||||
query[:,:mask_map.video_token_num],
|
||||
key[:,:key.shape[1], :, :],
|
||||
value[:,:key.shape[1], :, :],
|
||||
mask_id=None,
|
||||
is_causal=False,
|
||||
tensor_layout="NHD",
|
||||
)[0]
|
||||
|
||||
# if pre_defined_mask is not None:
|
||||
# output_text = flashinfer.single_prefill_with_kv_cache(
|
||||
# q=query[mask_map.video_token_num:, :, :],
|
||||
# k=key[:pre_defined_mask[0].sum(), :, :],
|
||||
# v=value[:pre_defined_mask[0].sum(), :, :],
|
||||
# causal=False,
|
||||
# return_lse=False,
|
||||
# )
|
||||
# return torch.cat([output_video, output_text], dim=0)
|
||||
# else:
|
||||
return output_video
|
||||
|
||||
return output_video.unsqueeze(0).contiguous()
|
||||
|
||||
# sparse-sageattention only supports (b, h, s, d) layout, need rearrange first
|
||||
query_hnd = rearrange(query.unsqueeze(0), "b s h d -> b h s d")
|
||||
key_hnd = rearrange(key.unsqueeze(0), "b s h d -> b h s d")
|
||||
value_hnd = rearrange(value.unsqueeze(0), "b s h d -> b h s d")
|
||||
converted_mask = repeat(sparge_mask_convert(mask=video_mask, block_size=block_size), "s t -> b h s t", b=query_hnd.shape[0], h=query_hnd.shape[1])
|
||||
converted_mask = repeat(sparge_mask_convert(mask=video_mask, block_size=block_size), "s t -> b h s t", b=1, h=query.shape[-2])
|
||||
|
||||
converted_mask = converted_mask.to(torch.int8)
|
||||
if pre_defined_mask is None:
|
||||
# wan case
|
||||
output = sparse_sageattn(
|
||||
query_hnd[:, :, :mask_map.video_token_num, :],
|
||||
key_hnd[:, :, :mask_map.video_token_num, :],
|
||||
value_hnd[:, :, :mask_map.video_token_num, :],
|
||||
mask_id=converted_mask,
|
||||
is_causal=False,
|
||||
tensor_layout="HND",
|
||||
)
|
||||
|
||||
# rearrange back to (s, h, d), we know that b = 1
|
||||
output = rearrange(output, "b h s d -> s (b h) d", b=1)
|
||||
return output
|
||||
|
||||
# query_video = query_hnd[:, :, :mask_map.video_token_num, :]
|
||||
# key_video = key_hnd
|
||||
# value_video = value_hnd
|
||||
# kv_border = (pre_defined_mask[0].sum() + 63) // 64
|
||||
# converted_mask[:, :, :, kv_border:] = False
|
||||
# output_video = sparse_sageattn(
|
||||
# query_video,
|
||||
# key_video,
|
||||
# value_video,
|
||||
# mask_id=converted_mask[:, :, :mask_map.video_token_num // block_size, :],
|
||||
# is_causal=False,
|
||||
# tensor_layout="HND",
|
||||
# )
|
||||
output = sparse_sageattn(
|
||||
query.transpose(1, 2)[:, :, :mask_map.video_token_num, :],
|
||||
key.transpose(1, 2)[:, :, :mask_map.video_token_num, :],
|
||||
value.transpose(1, 2)[:, :, :mask_map.video_token_num, :],
|
||||
mask_id=converted_mask,
|
||||
is_causal=False,
|
||||
tensor_layout="HND",
|
||||
)
|
||||
|
||||
# # rearrange back to (s, h, d), we know that b = 1
|
||||
# output_video = rearrange(output_video, "b h s d -> s (b h) d", b=1)
|
||||
|
||||
# # gt = sparse_sageattn(
|
||||
# # query_video,
|
||||
# # key_video,
|
||||
# # value_video,
|
||||
# # mask_id=None,
|
||||
# # is_causal=False,
|
||||
# # tensor_layout="HND",
|
||||
# # )[0]
|
||||
|
||||
|
||||
|
||||
# # import pdb; pdb.set_trace()
|
||||
|
||||
# output_text = flashinfer.single_prefill_with_kv_cache(
|
||||
# q=query[mask_map.video_token_num:, :, :],
|
||||
# k=key[:pre_defined_mask[0].sum(), :, :],
|
||||
# v=value[:pre_defined_mask[0].sum(), :, :],
|
||||
# causal=False,
|
||||
# return_lse=False,
|
||||
# )
|
||||
|
||||
# return torch.cat([output_video, output_text], dim=0)
|
||||
|
||||
def RadialAttention(query, key, value, mask_map=None, sparsity_type="radial", block_size=128, decay_factor=1, model_type=None, pre_defined_mask=None, use_sage_attention=False):
|
||||
#orig_seqlen, num_head, hidden_dim = query.shape
|
||||
return output.transpose(1, 2).contiguous()
|
||||
|
||||
def RadialAttention(query, key, value, mask_map=None, sparsity_type="radial", block_size=128, decay_factor=1, use_sage_attention=False):
|
||||
if sparsity_type == "dense":
|
||||
video_mask = torch.ones((mask_map.video_token_num // block_size, mask_map.video_token_num // block_size), device=query.device, dtype=torch.bool)
|
||||
video_mask = torch.ones((mask_map.video_token_num // block_size, mask_map.video_token_num // block_size), device=device, dtype=torch.bool)
|
||||
else:
|
||||
video_mask = mask_map.queryLogMask(query, sparsity_type, block_size=block_size, decay_factor=decay_factor, model_type=model_type) if mask_map else None
|
||||
video_mask = mask_map.queryLogMask(query.shape[0] * query.shape[1], sparsity_type, block_size=block_size, decay_factor=decay_factor) if mask_map else None
|
||||
|
||||
# backend = "sparse_sageattn" if use_sage_attention else "flashinfer"
|
||||
|
||||
# if backend == "flashinfer":
|
||||
# video_mask = video_mask[:mask_map.video_token_num // block_size, :mask_map.video_token_num // block_size]
|
||||
# # perform block-sparse attention on the video tokens
|
||||
# workspace_buffer = torch.empty(128 * 1024 * 1024, device=query.device, dtype=torch.uint8)
|
||||
# bsr_wrapper = flashinfer.BlockSparseAttentionWrapper(
|
||||
# workspace_buffer,
|
||||
# backend="fa2",
|
||||
# )
|
||||
|
||||
# indptr = get_indptr_from_mask(video_mask, query)
|
||||
# indices = get_indices_from_mask(video_mask, query)
|
||||
|
||||
# bsr_wrapper.plan(
|
||||
# indptr=indptr,
|
||||
# indices=indices,
|
||||
# M=mask_map.video_token_num,
|
||||
# N=mask_map.video_token_num,
|
||||
# R=block_size,
|
||||
# C=block_size,
|
||||
# num_qo_heads=num_head,
|
||||
# num_kv_heads=num_head,
|
||||
# head_dim=hidden_dim,
|
||||
# )
|
||||
|
||||
# return FlashInferBackend(query, key, value, mask_map, pre_defined_mask, bsr_wrapper)
|
||||
# elif backend == "sparse_sageattn":
|
||||
return SpargeSageAttnBackend(query, key, value, mask_map, video_mask, pre_defined_mask, block_size=block_size)
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# query = torch.randn(1, 2, 4, 64).cuda()
|
||||
# # mask = torch.tensor([
|
||||
# # [True, False, True, False],
|
||||
# # [False, True, False, True],
|
||||
# # [True, False, False, True],
|
||||
# # [False, True, True, False]
|
||||
# # ], dtype=torch.bool)
|
||||
# # indices = get_indices_from_mask(mask, query)
|
||||
# # indptr = get_indptr_from_mask(mask, query)
|
||||
# # print("Indices: ", indices)
|
||||
# # print("Indptr: ", indptr)
|
||||
# video_token_num = 3840 * 30
|
||||
# num_frame = 30
|
||||
# token_per_frame = video_token_num / num_frame
|
||||
# padded_video_token_num = ((video_token_num + 1) // 128 + 1) * 128
|
||||
# print("padded: ", padded_video_token_num)
|
||||
# temporal_mask = gen_log_mask_shrinked(query, padded_video_token_num, video_token_num, num_frame, sparse_type="radial", decay_factor=1, model_type="hunyuan")
|
||||
# plt.figure(figsize=(10, 8), dpi=500)
|
||||
|
||||
# plt.imshow(temporal_mask.cpu().numpy()[:, :], cmap='hot')
|
||||
# plt.colorbar()
|
||||
# plt.title("Temporal Mask")
|
||||
|
||||
# plt.savefig("temporal_mask.png",
|
||||
# dpi=300,
|
||||
# bbox_inches='tight',
|
||||
# pad_inches=0.1)
|
||||
|
||||
# plt.close()
|
||||
# # save the mask tensor
|
||||
# torch.save(temporal_mask, "temporal_mask.pt")
|
||||
return SpargeSageAttnBackend(query, key, value, mask_map, video_mask, block_size=block_size)
|
||||
Reference in New Issue
Block a user