stage

assets/strategy_convert.py

@@ -0,0 +1,21 @@
+import json
+
+initial_path = "/workspace/codefolder/FastVideo/assets/mask_strategy_hunyuan.json"
+
+with open(initial_path, "r") as f:
+    data = json.load(f)
+    for key in data.keys():
+        t, h, w = data[key]
+        data[key] = [min(2,t), h, w]
+
+save_path = "/workspace/codefolder/FastVideo/assets/test_mask_strategy_hunyuan.json"
+
+
+# print unique values for new data
+unique_values = set()
+for key in data.keys():
+    unique_values.add(tuple(data[key]))
+print(unique_values)
+
+with open(save_path, "w") as f:
+    json.dump(data, f, indent=4)

csrc/sliding_tile_attention/test/sba.py

@@ -0,0 +1,238 @@
+# def mask(b, h, q_idx, kv_idx):
+#     return kv_idx < text_length + img_seq_len
+from torch.nn.attention.flex_attention import create_block_mask, or_masks
+from torch import IntTensor, BoolTensor
+import torch
+from typing import Any, Callable, Dict, List, Optional, Union, Tuple
+import math
+# Peiyuan: This is neccesay. Dont know why. see https://github.com/pytorch/pytorch/issues/135028
+torch._inductor.config.realize_opcount_threshold = 100
+def generate_sba_mask(
+    canvas_twh, 
+    kernel_twh, 
+    tile_twh,
+    text_length
+):
+    """Generates a 3D NATTEN attention mask with a given kernel size.
+    
+    Args:
+        canvas_t: The time dimension of the canvas.
+        canvas_h: The height of the canvas.
+        canvas_w: The width of the canvas.
+        kernel_t: The time dimension of the kernel.
+        kernel_h: The height of the kernel.
+        kernel_w: The width of the kernel.
+    """
+    canvas_t, canvas_h, canvas_w = canvas_twh
+    kernel_t, kernel_h, kernel_w = kernel_twh
+    tile_t_size, tile_h_size, tile_w_size = tile_twh
+    total_tile_size = tile_t_size * tile_h_size * tile_w_size
+    canvas_tile_t, canvas_tile_h, canvas_tile_w = canvas_t // tile_t_size, canvas_h // tile_h_size, canvas_w // tile_w_size
+    img_seq_len = canvas_t * canvas_h * canvas_w
+    def get_tile_t_x_y(idx: IntTensor) -> Tuple[IntTensor, IntTensor, IntTensor]:
+        tile_id = idx // total_tile_size
+        tile_t = tile_id // (canvas_tile_h * canvas_tile_w)
+        tile_h = (tile_id % (canvas_tile_h * canvas_tile_w)) // canvas_tile_w
+        tile_w = tile_id % canvas_tile_w
+        return tile_t, tile_h, tile_w
+
+    def natten_mask_mod_3d(
+        b: IntTensor,
+        h: IntTensor,
+        q_idx: IntTensor,
+        kv_idx: IntTensor,
+    ) -> BoolTensor:
+        q_t_tile, q_x_tile, q_y_tile = get_tile_t_x_y(q_idx)
+        kv_t_tile, kv_x_tile, kv_y_tile = get_tile_t_x_y(kv_idx)
+        # kernel nominally attempts to center itself on the query, but kernel center
+        # is clamped to a fixed distance (kernel half-length) from the canvas edge
+        kernel_center_t = q_t_tile.clamp(kernel_t // 2, (canvas_tile_t - 1) - kernel_t // 2)
+        kernel_center_x = q_x_tile.clamp(kernel_h // 2, (canvas_tile_h - 1) - kernel_h // 2)
+        kernel_center_y = q_y_tile.clamp(kernel_w // 2, (canvas_tile_w - 1) - kernel_w // 2)
+        time_mask = (kernel_center_t - kv_t_tile).abs() <= kernel_t // 2
+        hori_mask = (kernel_center_x - kv_x_tile).abs() <= kernel_h // 2
+        vert_mask = (kernel_center_y - kv_y_tile).abs() <= kernel_w // 2
+        image_mask = (q_idx < img_seq_len) & (kv_idx < img_seq_len)
+        image_to_text_mask = (q_idx < img_seq_len) & (kv_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length)
+        text_to_all_mask = (q_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length)
+        return (image_mask & time_mask & hori_mask & vert_mask) | image_to_text_mask | text_to_all_mask
+
+    natten_mask_mod_3d.__name__ = f"natten_3d_c{canvas_t}x{canvas_w}x{canvas_h}_k{kernel_t}x{kernel_w}x{kernel_h}"
+    return natten_mask_mod_3d
+
+
+
+def generate_baseline_3d_window_mask(
+    canvas_twh, 
+    kernel_twh, 
+    img_seq_len,
+    text_length
+):
+    """Generates a 3D NATTEN attention mask with a given kernel size.
+    
+    Args:
+        canvas_t: The time dimension of the canvas.
+        canvas_h: The height of the canvas.
+        canvas_w: The width of the canvas.
+        kernel_t: The time dimension of the kernel.
+        kernel_h: The height of the kernel.
+        kernel_w: The width of the kernel.
+    """
+    canvas_t, canvas_h, canvas_w = canvas_twh
+    kernel_t, kernel_h, kernel_w = kernel_twh
+    def get_t_x_y(idx: IntTensor) -> Tuple[IntTensor, IntTensor, IntTensor]:
+        t = idx // (canvas_h * canvas_w)
+        x = (idx % (canvas_h * canvas_w)) // canvas_w
+        y = idx % canvas_w
+        return t, x, y
+
+    def natten_mask_mod_3d(
+        b: IntTensor,
+        h: IntTensor,
+        q_idx: IntTensor,
+        kv_idx: IntTensor,
+    ) -> BoolTensor:
+        q_t, q_x, q_y = get_t_x_y(q_idx)
+        kv_t, kv_x, kv_y = get_t_x_y(kv_idx)
+        # kernel nominally attempts to center itself on the query, but kernel center
+        # is clamped to a fixed distance (kernel half-length) from the canvas edge
+        kernel_center_t = q_t.clamp(kernel_t // 2, (canvas_t - 1) - kernel_t // 2)
+        kernel_center_x = q_x.clamp(kernel_h // 2, (canvas_h - 1) - kernel_h // 2)
+        kernel_center_y = q_y.clamp(kernel_w // 2, (canvas_w - 1) - kernel_w // 2)
+        time_mask = (kernel_center_t - kv_t).abs() <= kernel_t // 2
+        hori_mask = (kernel_center_x - kv_x).abs() <= kernel_h // 2
+        vert_mask = (kernel_center_y - kv_y).abs() <= kernel_w // 2
+        no_pad_mask = kv_idx < text_length + img_seq_len
+        return time_mask & hori_mask & vert_mask & no_pad_mask
+
+    natten_mask_mod_3d.__name__ = f"natten_3d_c{canvas_t}x{canvas_w}x{canvas_h}_k{kernel_t}x{kernel_w}x{kernel_h}"
+    return natten_mask_mod_3d
+
+
+def generate_tiled_3d_window_mask(
+    canvas_twh, 
+    kernel_twh, 
+    img_seq_len,
+    text_length
+):
+    """Generates a 3D NATTEN attention mask with a given kernel size.
+    
+    Args:
+        canvas_t: The time dimension of the canvas.
+        canvas_h: The height of the canvas.
+        canvas_w: The width of the canvas.
+        kernel_t: The time dimension of the kernel.
+        kernel_h: The height of the kernel.
+        kernel_w: The width of the kernel.
+    """
+    canvas_t, canvas_h, canvas_w = canvas_twh
+    kernel_t, kernel_h, kernel_w = kernel_twh
+    tile_t, tile_h, tile_w = 4, 8, 8
+    n_tile_t, n_tile_h, n_tile_w = canvas_t // tile_t, canvas_h // tile_h, canvas_w // tile_w
+    
+    def get_t_x_y(idx: IntTensor) -> Tuple[IntTensor, IntTensor, IntTensor]:
+        tile_id = idx // (tile_t * tile_h * tile_w)
+        t_t, t_x, t_y = tile_id // (n_tile_h * n_tile_w), (tile_id % (n_tile_h * n_tile_w)) // n_tile_w, tile_id % n_tile_w
+        t_offset = idx % (tile_t * tile_h * tile_w)
+        i_t, i_x, i_y = t_offset // (tile_h * tile_w), (t_offset % (tile_h * tile_w)) // tile_w, t_offset % tile_w
+        return t_t * tile_t + i_t, t_x * tile_h + i_x, t_y * tile_w + i_y
+    
+    def natten_mask_mod_3d(
+        b: IntTensor,
+        h: IntTensor,
+        q_idx: IntTensor,
+        kv_idx: IntTensor,
+    ) -> BoolTensor:
+        q_t, q_x, q_y = get_t_x_y(q_idx)
+        kv_t, kv_x, kv_y = get_t_x_y(kv_idx)
+        # kernel nominally attempts to center itself on the query, but kernel center
+        # is clamped to a fixed distance (kernel half-length) from the canvas edge
+        kernel_center_t = q_t.clamp(kernel_t // 2, (canvas_t - 1) - kernel_t // 2)
+        kernel_center_x = q_x.clamp(kernel_h // 2, (canvas_h - 1) - kernel_h // 2)
+        kernel_center_y = q_y.clamp(kernel_w // 2, (canvas_w - 1) - kernel_w // 2)
+        time_mask = (kernel_center_t - kv_t).abs() <= kernel_t // 2
+        hori_mask = (kernel_center_x - kv_x).abs() <= kernel_h // 2
+        vert_mask = (kernel_center_y - kv_y).abs() <= kernel_w // 2
+        no_pad_mask = kv_idx < text_length + img_seq_len
+        return time_mask & hori_mask & vert_mask & no_pad_mask
+
+    natten_mask_mod_3d.__name__ = f"natten_3d_c{canvas_t}x{canvas_w}x{canvas_h}_k{kernel_t}x{kernel_w}x{kernel_h}"
+    return natten_mask_mod_3d
+
+
+def generate_text_mask(img_seq_len, text_length):
+    def text_mask(b, h, q_idx, kv_idx):
+        mask1 = kv_idx < text_length + img_seq_len
+        mask2 = kv_idx >= img_seq_len
+        return mask1 & mask2
+    return text_mask
+
+ 
+def get_sliding_block_attention_mask(kernel_size, tile_size, img_size, text_length, device):
+    img_seq_len = img_size[0] * img_size[1] * img_size[2]
+    image_mask = generate_sba_mask(img_size, kernel_size, tile_size, text_length)
+    mask = create_block_mask(image_mask, B=None, H=None, Q_LEN=img_seq_len+256 , KV_LEN=img_seq_len+256, device=device,  _compile=True)
+    return mask
+
+def get_baseline_sliding_window_mask(kernel_size, img_seq_len, text_length, device):
+    image_mask = generate_baseline_3d_window_mask( (32, 48, 80), kernel_size, img_seq_len, text_length)
+    text_mask = generate_text_mask(img_seq_len=img_seq_len, text_length=text_length)
+    mask = or_masks(image_mask, text_mask)
+    mask = create_block_mask(mask, B=None, H=None, Q_LEN=img_seq_len + 256 , KV_LEN=img_seq_len + 256, device=device,  _compile=True)
+    return mask
+
+def get_tiled_sliding_window_mask(kernel_size, img_seq_len, text_length, device):
+    image_mask = generate_tiled_3d_window_mask((32, 48, 80), kernel_size, img_seq_len, text_length)
+    text_mask = generate_text_mask(img_seq_len=img_seq_len, text_length=text_length)
+    mask = or_masks(image_mask, text_mask)
+    mask = create_block_mask(mask, B=None, H=None, Q_LEN=img_seq_len + 256 , KV_LEN=img_seq_len + 256, device=device,  _compile=True)
+    return mask
+
+
+def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True):
+    seq_length = q_all.shape[2]
+    # if has_text:
+    #     assert q_all.shape[
+    #         2] == 115456, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
+    #     assert q_all.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
+    #     target_size = math.ceil(seq_length / 384) * 384
+    #     pad_size = target_size - seq_length
+    #     if pad_size > 0:
+    #         q_all = torch.cat([q_all, q_all[:, :, -pad_size:]], dim=2)
+    #         k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
+    #         v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
+    # else:
+    #     assert q_all.shape[2] == 82944
+
+    hidden_states = torch.empty_like(q_all)
+    # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
+    for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
+        for batch in range(q_all.shape[0]):
+            q_head, k_head, v_head, o_head = (q_all[batch:batch + 1, head_index:head_index + 1],
+                                              k_all[batch:batch + 1,
+                                                    head_index:head_index + 1], v_all[batch:batch + 1,
+                                                                                      head_index:head_index + 1],
+                                              hidden_states[batch:batch + 1, head_index:head_index + 1])
+
+            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text)
+    if has_text:
+        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True)
+    return hidden_states[:, :, :seq_length]
+
+if __name__ == "__main__":
+    # benchmark speed 
+    from torch.nn.attention.flex_attention import flex_attention
+    flex_attention = torch.compile(flex_attention)
+    import time
+    device = torch.device("cuda")
+    kernel_size_ls = [(8, 6, 10), (4, 6, 10), (4, 6, 5), (4, 3, 5), (3, 3, 3)]
+    tile_size = (4, 8, 8)
+    
+    # random input
+    q = torch.randn(1, 24, 123136, 128, device=device, dtype=torch.bfloat16)
+    k = torch.randn(1, 24, 123136, 128, device=device, dtype=torch.bfloat16)
+    v = torch.randn(1, 24, 123136, 128, device=device, dtype=torch.bfloat16)
+    
+    for kernel_size in kernel_size_ls:
+        mask = get_sliding_block_attention_mask(kernel_size, tile_size, (32, 48, 80), 39, device)
+        

fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py

@@ -813,7 +813,19 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                 t, l, h = map(int, key.split('_'))
                 result[t][l][h] = value
             return result
-
+        
+        selected_strategies = [(2, 6, 1), (1, 6, 10), (2, 3, 3), (2, 6, 10), (2, 1, 10), (2, 3, 5)]
+        text_length = prompt_mask.sum()
+        selected_attn_processor = []
+        from torch.nn.attention.flex_attention import flex_attention
+        from functools import partial
+        from csrc.sliding_tile_attention.test.sba import get_sliding_block_attention_mask
+        
+        for ms in selected_strategies:
+            mask = get_sliding_block_attention_mask(ms, (6, 8, 8), (12, 48, 80), text_length, self.transformer.device)
+            attn_processor = torch.compile(partial(flex_attention, block_mask=mask))
+            selected_attn_processor.append(attn_processor)
+        
         mask_strategy = dict_to_3d_list(mask_strategy)
         # if is_progress_bar:
         with self.progress_bar(total=num_inference_steps) as progress_bar:
@@ -849,6 +861,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                         mask_strategy=mask_strategy[i],
                         guidance=guidance_expand,
                         return_dict=False,
+                        selected_attn_processor=selected_attn_processor,
                     )[0]
 
                 # perform guidance
@@ -904,14 +917,20 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                 latents = (latents / self.vae.config.scaling_factor + self.vae.config.shift_factor)
             else:
                 latents = latents / self.vae.config.scaling_factor
+                
+            self.transformer = self.transformer.to('cpu')
 
             with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
                 if enable_tiling:
+                    print("tiling is enabled")
                     self.vae.enable_tiling()
                 if enable_vae_sp:
                     self.vae.enable_parallel()
+                    
                 image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
 
+            self.transformer = self.transformer.to(device)
+            
             if expand_temporal_dim or image.shape[2] == 1:
                 image = image.squeeze(2)
 

fastvideo/models/hunyuan/inference.py

@@ -15,7 +15,7 @@ from fastvideo.models.hunyuan.text_encoder import TextEncoder
 from fastvideo.models.hunyuan.utils.data_utils import align_to
 from fastvideo.models.hunyuan.vae import load_vae
 from fastvideo.utils.parallel_states import nccl_info
-
+from fastvideo.models.hunyuan.modules.fp8 import convert_fp8_linear
 
 class Inference(object):
 
@@ -76,7 +76,7 @@ class Inference(object):
 
         # =========================== Build main model ===========================
         logger.info("Building model...")
-        factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]}
+        factor_kwargs = {"device": 'cpu', "dtype": PRECISION_TO_TYPE[args.precision]}
         in_channels = args.latent_channels
         out_channels = args.latent_channels
 
@@ -86,6 +86,11 @@ class Inference(object):
             out_channels=out_channels,
             factor_kwargs=factor_kwargs,
         )
+        
+        if args.use_fp8:
+            print("loading fp8 model")
+            convert_fp8_linear(model, args.dit_weight, original_dtype=PRECISION_TO_TYPE[args.precision])
+        
         model = model.to(device)
         model = Inference.load_state_dict(args, model, pretrained_model_path)
         if args.enable_torch_compile:

fastvideo/models/hunyuan/modules/attenion.py

@@ -1,18 +1,22 @@
 import torch
 import torch.nn.functional as F
 from einops import rearrange
+import time
+# try:
+#     from st_attn import sliding_tile_attention
+# except ImportError:
+#     print("Could not load Sliding Tile Attention.")
+#     sliding_tile_attention = None
 
-try:
-    from st_attn import sliding_tile_attention
-except ImportError:
-    print("Could not load Sliding Tile Attention.")
-    sliding_tile_attention = None
+import tk_4090_cuda as tk
 
+from csrc.sliding_tile_attention.test.sba import sliding_tile_attention
+from torch.nn.attention.flex_attention import flex_attention
 from fastvideo.models.flash_attn_no_pad import flash_attn_no_pad
 from fastvideo.utils.communications import all_gather, all_to_all_4D
 from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
 
-
+flex_attention = torch.compile(flex_attention, dynamic=False)
 def attention(
     q,
     k,
@@ -35,10 +39,10 @@ def attention(
 
 
 def tile(x, sp_size):
-    x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80)
+    x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=12 // sp_size, h=48, w=80)
     return rearrange(x,
                      "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-                     n_t=5,
+                     n_t=2,
                      n_h=6,
                      n_w=10,
                      ts_t=6,
@@ -49,16 +53,16 @@ def tile(x, sp_size):
 def untile(x, sp_size):
     x = rearrange(x,
                   "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-                  n_t=5,
+                  n_t=2,
                   n_h=6,
                   n_w=10,
                   ts_t=6,
                   ts_h=8,
                   ts_w=8)
-    return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80)
+    return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=12 // sp_size, h=48, w=80)
 
 
-def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=None):
+def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=None, selected_attn_processor=None):
     query, encoder_query = q
     key, encoder_key = k
     value, encoder_value = v
@@ -91,21 +95,88 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=
         current_rank = nccl_info.rank_within_group
         start_head = current_rank * head_num
         windows = [mask_strategy[head_idx + start_head] for head_idx in range(head_num)]
-
-        hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2)
+        
+        hidden_states = torch.empty_like(query)
+        
+        time_start = time.time()
+        torch.cuda.synchronize()
+        
+        # process the qkv head by head
+        for head_index in range(head_num):
+            query_head = query[:, head_index:head_index + 1, :, :]
+            key_head = key[:, head_index:head_index + 1, :, :]
+            value_head = value[:, head_index:head_index + 1, :, :]
+            mask = windows[head_index]
+            
+            # selected_attn_processor [(2, 6, 1), (1, 6, 10), (2, 3, 3), (2, 6, 10), (2, 1, 10), (2, 3, 5)]
+            
+            # print(f"mask_strategy: {mask_strategy[head_index]}")
+            
+            if mask_strategy[head_index] == [2, 6, 1]:
+                head_hidden_states = selected_attn_processor[0](query_head, key_head, value_head).unsqueeze(0).transpose(1, 2)
+            elif mask_strategy[head_index] == [1, 6, 10]:
+                head_hidden_states = selected_attn_processor[1](query_head, key_head, value_head).unsqueeze(0).transpose(1, 2)
+            elif mask_strategy[head_index] == [2, 3, 3]:
+                head_hidden_states = selected_attn_processor[2](query_head, key_head, value_head).unsqueeze(0).transpose(1, 2)
+            elif mask_strategy[head_index] == [2, 1, 10]:
+                head_hidden_states = selected_attn_processor[4](query_head, key_head, value_head).unsqueeze(0).transpose(1, 2)
+            elif mask_strategy[head_index] == [2, 3, 5]:
+                head_hidden_states = selected_attn_processor[5](query_head, key_head, value_head).unsqueeze(0).transpose(1, 2)
+            else:
+                head_hidden_states = selected_attn_processor[3](query_head, key_head, value_head).unsqueeze(0).transpose(1, 2)
+                
+            hidden_states[:, head_index:head_index + 1, :, :] = head_hidden_states
+        torch.cuda.synchronize()
+        time_end = time.time()
+        print(f"Time taken for sliding tile attention: {time_end - time_start}")
     else:
-        query = torch.cat([query, encoder_query], dim=1)
-        key = torch.cat([key, encoder_key], dim=1)
-        value = torch.cat([value, encoder_value], dim=1)
+        # query = torch.cat([query, encoder_query], dim=1)
+        # key = torch.cat([key, encoder_key], dim=1)
+        # value = torch.cat([value, encoder_value], dim=1)
         # B, S, 3, H, D
-        qkv = torch.stack([query, key, value], dim=2)
+        # result = torch.empty_like(query)
+        # start_time = time.time()
+        # torch.cuda.synchronize()
 
-        attn_mask = F.pad(text_mask, (sequence_length, 0), value=True)
-        hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None)
+        # output = tk.attention_fwd_4090(query, key, value, result, text_length)
+        
+        # torch.cuda.synchronize()
+        # end_time = time.time()
+        # print(f"Time taken for tk attention: {end_time - start_time}")
+        
+        # qkv = torch.stack([query, key, value], dim=2)
+
+        # attn_mask = F.pad(text_mask, (sequence_length, 0), value=True)
+        
+
+        # flash_start_time = time.time()
+        # torch.cuda.synchronize()        
+        # hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None)
+        # torch.cuda.synchronize()
+        # flash_end_time = time.time()
+        # print(f"Time taken for flash attention: {flash_end_time - flash_start_time}")
+        
+        
+        flex_start_time = time.time()
+        query = torch.cat([tile(query, nccl_info.sp_size), encoder_query], dim=1).transpose(1, 2)
+        key = torch.cat([tile(key, nccl_info.sp_size), encoder_key], dim=1).transpose(1, 2)
+        value = torch.cat([tile(value, nccl_info.sp_size), encoder_value], dim=1).transpose(1, 2)
+        torch.cuda.synchronize()
+        mask = get_sliding_block_attention_mask((2,3,5), (6, 8, 8), (12, 48, 80), text_length, 'cuda')
+        # print(f"query shape: {query.shape}")
+        # print(f"key shape: {key.shape}")
+        # print(f"value shape: {value.shape}")
+        
+        flex_hidden_states = flex_attention(query, key, value, block_mask=mask).transpose(1, 2)
+        torch.cuda.synchronize()
+        flex_end_time = time.time()
+        print(f"Time taken for flex attention: {flex_end_time - flex_start_time}")
+        
+        
+        hidden_states = flex_hidden_states
 
     hidden_states, encoder_hidden_states = hidden_states.split_with_sizes((sequence_length, encoder_sequence_length),
                                                                           dim=1)
-
     if mask_strategy[0] is not None:
         hidden_states = untile(hidden_states, nccl_info.sp_size)
 

fastvideo/models/hunyuan/modules/fp8.py

@@ -0,0 +1,101 @@
+import os
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+def get_fp_maxval(bits=8, mantissa_bit=3, sign_bits=1):
+    _bits = torch.tensor(bits)
+    _mantissa_bit = torch.tensor(mantissa_bit)
+    _sign_bits = torch.tensor(sign_bits)
+    M = torch.clamp(torch.round(_mantissa_bit), 1, _bits - _sign_bits)
+    E = _bits - _sign_bits - M
+    bias = 2 ** (E - 1) - 1
+    mantissa = 1
+    for i in range(mantissa_bit - 1):
+        mantissa += 1 / (2 ** (i+1))
+    maxval = mantissa * 2 ** (2**E - 1 - bias)
+    return maxval
+
+def quantize_to_fp8(x, bits=8, mantissa_bit=3, sign_bits=1):
+    """
+    Default is E4M3.
+    """
+    bits = torch.tensor(bits)
+    mantissa_bit = torch.tensor(mantissa_bit)
+    sign_bits = torch.tensor(sign_bits)
+    M = torch.clamp(torch.round(mantissa_bit), 1, bits - sign_bits)
+    E = bits - sign_bits - M
+    bias = 2 ** (E - 1) - 1
+    mantissa = 1
+    for i in range(mantissa_bit - 1):
+        mantissa += 1 / (2 ** (i+1))
+    maxval = mantissa * 2 ** (2**E - 1 - bias)
+    minval = - maxval
+    minval = - maxval if sign_bits == 1 else torch.zeros_like(maxval)
+    input_clamp = torch.min(torch.max(x, minval), maxval)
+    log_scales = torch.clamp((torch.floor(torch.log2(torch.abs(input_clamp)) + bias)).detach(), 1.0)
+    log_scales = 2.0 ** (log_scales - M - bias.type(x.dtype))
+    # dequant
+    qdq_out = torch.round(input_clamp / log_scales) * log_scales
+    return qdq_out, log_scales
+
+def fp8_tensor_quant(x, scale, bits=8, mantissa_bit=3, sign_bits=1):
+    for i in range(len(x.shape) - 1):
+        scale = scale.unsqueeze(-1)
+    new_x = x / scale
+    quant_dequant_x, log_scales = quantize_to_fp8(new_x, bits=bits, mantissa_bit=mantissa_bit, sign_bits=sign_bits)
+    return quant_dequant_x, scale, log_scales
+
+def fp8_activation_dequant(qdq_out, scale, dtype):
+    qdq_out = qdq_out.type(dtype)
+    quant_dequant_x = qdq_out * scale.to(dtype)
+    return quant_dequant_x
+
+def fp8_linear_forward(cls, original_dtype, input):
+    weight_dtype = cls.weight.dtype
+    #####
+    if cls.weight.dtype != torch.float8_e4m3fn:
+        maxval = get_fp_maxval()
+        scale = torch.max(torch.abs(cls.weight.flatten())) / maxval
+        linear_weight, scale, log_scales = fp8_tensor_quant(cls.weight, scale)
+        linear_weight = linear_weight.to(torch.float8_e4m3fn)
+        weight_dtype = linear_weight.dtype
+    else:
+        scale = cls.fp8_scale.to(cls.weight.device)
+        linear_weight = cls.weight
+    #####
+
+    if weight_dtype == torch.float8_e4m3fn and cls.weight.sum() != 0:
+        if True or len(input.shape) == 3:
+            cls_dequant = fp8_activation_dequant(linear_weight, scale, original_dtype)
+            if cls.bias != None:
+                output = F.linear(input, cls_dequant, cls.bias)
+            else:
+                output = F.linear(input, cls_dequant)
+            return output
+        else:
+            return cls.original_forward(input.to(original_dtype))
+    else:
+        return cls.original_forward(input)
+
+def convert_fp8_linear(module, dit_weight_path, original_dtype, params_to_keep={}):
+    setattr(module, "fp8_matmul_enabled", True)
+
+    # loading fp8 mapping file
+    fp8_map_path = dit_weight_path.replace('.pt', '_map.pt')
+    if os.path.exists(fp8_map_path):
+        fp8_map = torch.load(fp8_map_path, map_location=lambda storage, loc: storage)
+    else:
+        raise ValueError(f"Invalid fp8_map path: {fp8_map_path}.")
+
+    fp8_layers = []
+    for key, layer in module.named_modules():
+        if isinstance(layer, nn.Linear) and ('double_blocks' in key or 'single_blocks' in key):
+            fp8_layers.append(key)
+            original_forward = layer.forward
+            layer.weight = torch.nn.Parameter(layer.weight.to(torch.float8_e4m3fn))
+            setattr(layer, "fp8_scale", fp8_map[key].to(dtype=original_dtype))
+            setattr(layer, "original_forward", original_forward)
+            setattr(layer, "forward", lambda input, m=layer: fp8_linear_forward(m, original_dtype, input))
+    

fastvideo/models/hunyuan/modules/models.py

@@ -110,6 +110,7 @@ class MMDoubleStreamBlock(nn.Module):
         freqs_cis: tuple = None,
         text_mask: torch.Tensor = None,
         mask_strategy=None,
+        selected_attn_processor=None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         (
             img_mod1_shift,
@@ -171,6 +172,7 @@ class MMDoubleStreamBlock(nn.Module):
             img_kv_len=img_k.shape[1],
             text_mask=text_mask,
             mask_strategy=mask_strategy,
+            selected_attn_processor=selected_attn_processor,
         )
 
         # attention computation end
@@ -260,6 +262,7 @@ class MMSingleStreamBlock(nn.Module):
         freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
         text_mask: torch.Tensor = None,
         mask_strategy=None,
+        selected_attn_processor=None,
     ) -> torch.Tensor:
         mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
         x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale)
@@ -296,6 +299,7 @@ class MMSingleStreamBlock(nn.Module):
             img_kv_len=img_k.shape[1],
             text_mask=text_mask,
             mask_strategy=mask_strategy,
+            selected_attn_processor=selected_attn_processor,
         )
 
         # attention computation end
@@ -520,6 +524,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = False,
         guidance=None,
+        selected_attn_processor=None,
     ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
         if guidance is None:
             guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16)
@@ -566,7 +571,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
         # --------------------- Pass through DiT blocks ------------------------
 
         for index, block in enumerate(self.double_blocks):
-            double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]]
+            double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index], selected_attn_processor]
             img, txt = block(*double_block_args)
         # Merge txt and img to pass through single stream blocks.
         x = torch.cat((img, txt), 1)
@@ -581,6 +586,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
                     (freqs_cos, freqs_sin),
                     text_mask,
                     mask_strategy[index + len(self.double_blocks)],
+                    selected_attn_processor,
                 ]
                 x = block(*single_block_args)
                 if output_features and _ % output_features_stride == 0:

fastvideo/sample/sample_t2v_hunyuan.py

@@ -1,7 +1,7 @@
 import argparse
 import os
 from pathlib import Path
-
+import json
 import imageio
 import numpy as np
 import torch
@@ -28,6 +28,8 @@ def main(args):
 
     print(args)
     models_root_path = Path(args.model_path)
+    with open(args.mask_strategy_file_path, 'r') as f:
+        mask_strategy = json.load(f)
     if not models_root_path.exists():
         raise ValueError(f"`models_root` not exists: {models_root_path}")
 
@@ -61,6 +63,7 @@ def main(args):
             flow_shift=args.flow_shift,
             batch_size=args.batch_size,
             embedded_guidance_scale=args.embedded_cfg_scale,
+            mask_strategy=mask_strategy,
         )
         videos = rearrange(outputs["samples"], "b c t h w -> t b c h w")
         outputs = []
@@ -162,6 +165,7 @@ if __name__ == "__main__":
     )
 
     # Model parameters
+    parser.add_argument("--use-fp8", action='store_true')
     parser.add_argument("--model", type=str, default="HYVideo-T/2-cfgdistill")
     parser.add_argument("--latent-channels", type=int, default=16)
     parser.add_argument("--precision", type=str, default="bf16", choices=["fp32", "fp16", "bf16"])
@@ -202,7 +206,8 @@ if __name__ == "__main__":
     parser.add_argument("--text-states-dim-2", type=int, default=768)
     parser.add_argument("--tokenizer-2", type=str, default="clipL")
     parser.add_argument("--text-len-2", type=int, default=77)
-
+    parser.add_argument("--vae_tiling", action='store_true')
+    parser.add_argument("--mask_strategy_file_path", type=str, default="assets/mask_strategy.json")
     args = parser.parse_args()
     # process for vae sequence parallel
     if args.vae_sp and not args.vae_tiling:

scripts/inference/inference_hunyuan.sh

@@ -1,32 +1,33 @@
 #!/bin/bash
 
-num_gpus=4
-export MODEL_BASE=data/FastHunyuan
-torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
-    fastvideo/sample/sample_t2v_hunyuan.py \
-    --height 720 \
-    --width 1280 \
-    --num_frames 125 \
-    --num_inference_steps 6 \
-    --guidance_scale 1 \
-    --embedded_cfg_scale 6 \
-    --flow_shift 17 \
-    --flow-reverse \
-    --prompt ./assets/prompt.txt \
-    --seed 1024 \
-    --output_path outputs_video/hunyuan/vae_sp/ \
-    --model_path $MODEL_BASE \
-    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
-    --vae-sp
+# num_gpus=4
+# export MODEL_BASE=data/FastHunyuan
+# torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+#     fastvideo/sample/sample_t2v_hunyuan.py \
+#     --height 720 \
+#     --width 1280 \
+#     --num_frames 125 \
+#     --num_inference_steps 6 \
+#     --guidance_scale 1 \
+#     --embedded_cfg_scale 6 \
+#     --flow_shift 17 \
+#     --flow-reverse \
+#     --prompt ./assets/prompt.txt \
+#     --seed 1024 \
+#     --output_path outputs_video/hunyuan/vae_sp/ \
+#     --model_path $MODEL_BASE \
+#     --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+#     --vae-sp
 
 # Inference original Hunyuan
-num_gpus=4
+num_gpus=1
 export MODEL_BASE=data/hunyuan
+mask_strategy_file_path=assets/test_mask_strategy_hunyuan.json
 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     fastvideo/sample/sample_t2v_hunyuan.py \
-    --height 720 \
+    --height 768 \
     --width 1280 \
-    --num_frames 125 \
+    --num_frames 45 \
     --num_inference_steps 50 \
     --guidance_scale 1 \
     --embedded_cfg_scale 6 \
@@ -34,7 +35,11 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     --flow-reverse \
     --prompt ./assets/prompt.txt \
     --seed 1024 \
-    --output_path outputs_video/hunyuan/vae_sp/ \
+    --output_path outputs_video/hunyuan/flex_attention_235/ \
     --model_path $MODEL_BASE \
-    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
-    --vae-sp
+    --mask_strategy_file_path $mask_strategy_file_path \
+    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt \
+    --vae-sp \
+    --use-fp8 \
+    --use-cpu-offload \
+    --vae_tiling

scripts/inference/inference_hunyuan_STA.sh

@@ -1,39 +1,39 @@
 #!/bin/bash
 
 # Inference with STA + Teacache
-num_gpus=1 # currently it should be a factor of 30
-mask_strategy_file_path=assets/mask_strategy_hunyuan.json
-export MODEL_BASE=data/hunyuan
-rel_l1_thresh=0.15
-CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
-    fastvideo/sample/sample_t2v_hunyuan_STA.py \
-    --height 768 \
-    --width 1280 \
-    --num_frames 117 \
-    --num_inference_steps 50 \
-    --guidance_scale 1 \
-    --embedded_cfg_scale 6 \
-    --flow_shift 7 \
-    --flow-reverse \
-    --prompt ./assets/prompt.txt \
-    --seed 12345 \
-    --output_path outputs_video/hunyuan_STA/ \
-    --model_path $MODEL_BASE \
-    --mask_strategy_file_path $mask_strategy_file_path \
-    --rel_l1_thresh $rel_l1_thresh \
-    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
-    --vae-sp \
-    --enable_teacache 
+# num_gpus=1 # currently it should be a factor of 30
+# mask_strategy_file_path=assets/mask_strategy_hunyuan.json
+# export MODEL_BASE=data/hunyuan
+# rel_l1_thresh=0.15
+# CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
+#     fastvideo/sample/sample_t2v_hunyuan_STA.py \
+#     --height 768 \
+#     --width 1280 \
+#     --num_frames 117 \
+#     --num_inference_steps 50 \
+#     --guidance_scale 1 \
+#     --embedded_cfg_scale 6 \
+#     --flow_shift 7 \
+#     --flow-reverse \
+#     --prompt ./assets/prompt.txt \
+#     --seed 12345 \
+#     --output_path outputs_video/hunyuan_STA/ \
+#     --model_path $MODEL_BASE \
+#     --mask_strategy_file_path $mask_strategy_file_path \
+#     --rel_l1_thresh $rel_l1_thresh \
+#     --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+#     --vae-sp \
+#     --enable_teacache 
 
 # Inference with STA only
 num_gpus=1
-mask_strategy_file_path=assets/mask_strategy_hunyuan.json
+mask_strategy_file_path=assets/test_mask_strategy_hunyuan.json
 export MODEL_BASE=data/hunyuan
-CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
     fastvideo/sample/sample_t2v_hunyuan_STA.py \
     --height 768 \
     --width 1280 \
-    --num_frames 117 \
+    --num_frames 45 \
     --num_inference_steps 50 \
     --guidance_scale 1 \
     --embedded_cfg_scale 6 \
@@ -41,9 +41,11 @@ CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_p
     --flow-reverse \
     --prompt ./assets/prompt.txt \
     --seed 12345 \
-    --output_path outputs_video/hunyuan_STA/ \
+    --output_path outputs_video/hunyuan_STA/test_4090/ \
     --model_path $MODEL_BASE \
     --mask_strategy_file_path $mask_strategy_file_path \
-    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt \
     --vae-sp \
-    --enable_torch_compile
+    --use-fp8 \
+    --use-cpu-offload \
+    --vae_tiling

scripts/inference/inference_hunyuan_hf.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-num_gpus=4
+num_gpus=1
 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     fastvideo/sample/sample_t2v_hunyuan_hf.py \
     --model_path ~/data/hunyuan_diffusers/ \

test.py

@@ -0,0 +1,58 @@
+import torch
+import thunderkittens_cuda as tk
+import time
+
+def test_attention_kernel():
+    # Create small test tensors
+    batch_size = 1
+    seq_len = 1  # Start small
+    n_heads = 4
+    head_dim = 64  # Must match your kernel's supported dimensions (64 or 128)
+    
+    # Create random inputs
+    q = torch.randn(batch_size, seq_len, n_heads, head_dim, 
+                   device='cuda', dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    v = torch.randn_like(q)
+    
+    print("Input shapes:", q.shape)
+    print("Input tensors created")
+    
+    o = torch.empty_like(q)
+    try:
+        # Add timing
+        start = time.time()
+        print("Calling thunderkittens kernel...")
+        _ = tk.attention_fwd_4090(q, k, v, o, False)
+        print("yes")
+        torch.cuda.synchronize()
+        end = time.time()
+        print(f"Kernel completed in {(end-start)*1000:.2f} ms")
+        print("Output shape:", o.shape)
+        
+        # Run again with different parameters
+        _ = tk.attention_fwd_4090(q, k, v, o, True)  # Try with causal=True
+        print("Causal version completed")
+        
+        # Validate results (compare with PyTorch implementation)
+        print("Computing reference result...")
+        q_scaled = q * (1.0 / torch.sqrt(torch.tensor(head_dim, dtype=torch.float)))
+        attn = torch.matmul(q_scaled, k.transpose(-1, -2))
+        attn = torch.softmax(attn, dim=-1)
+        ref_out = torch.matmul(attn, v)
+        
+        # Check if results are reasonable
+        error = (o - ref_out).abs().mean().item()
+        print(f"Mean absolute error: {error}")
+        
+        return True
+    except Exception as e:
+        print(f"Error calling kernel: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+if __name__ == "__main__":
+    print("Testing attention kernel...")
+    success = test_attention_kernel()
+    print(f"Test {'succeeded' if success else 'failed'}")