diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py index 5c74303..e6d867b 100644 --- a/wanvideo/modules/model.py +++ b/wanvideo/modules/model.py @@ -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, diff --git a/wanvideo/radial_attention/attn_mask.py b/wanvideo/radial_attention/attn_mask.py index e47c5ae..40384b4 100644 --- a/wanvideo/radial_attention/attn_mask.py +++ b/wanvideo/radial_attention/attn_mask.py @@ -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") \ No newline at end of file + return SpargeSageAttnBackend(query, key, value, mask_map, video_mask, block_size=block_size) \ No newline at end of file