syn with main

Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>

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

@@ -813,7 +813,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                 t, l, h = map(int, key.split('_'))
                 result[t][l][h] = value
             return result
-        torch._dynamo.config.cache_size_limit = 128
+
         mask_strategy = dict_to_3d_list(mask_strategy)
         # if is_progress_bar:
         with self.progress_bar(total=num_inference_steps) as progress_bar:

fastvideo/models/hunyuan/modules/attenion.py

@@ -6,68 +6,12 @@ try:
     from st_attn import sliding_tile_attention
 except ImportError:
     print("Could not load Sliding Tile Attention.")
-sliding_tile_attention = None
-from functools import lru_cache
-from torch.nn.attention.flex_attention import flex_attention
+    sliding_tile_attention = None
+
 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
-from csrc.sliding_tile_attention.test.flex_sta_ref import get_sliding_tile_attention_mask
-@lru_cache(maxsize=32)
-def get_compiled_flex_attention(strategy, tile_size, image_size, text_length, device):
-    """
-    Create and compile flex attention with a specific sliding block mask.
-    This function is cached to avoid recompiling for the same parameters.
-    
-    Args:
-        strategy (tuple): A tuple (t, h, w) defining the strategy
-        tile_size (tuple): A tuple (ts_t, ts_h, ts_w) defining the tile size
-        image_size (tuple): A tuple (n_t, n_h, n_w) defining the image size
-        text_length (int): The text length
-        device (str): The device to use
-        
-    Returns:
-        function: A compiled flex attention function with the specified mask
-    """
-    # Convert strategy to the required format (ceil(t*3/2), h*2, w)
-    adjusted_strategy = strategy
-    
-    # Get the sliding block attention mask
-    mask = get_sliding_tile_attention_mask(
-        adjusted_strategy, 
-        tile_size, 
-        image_size, 
-        text_length, 
-        device
-    )
-    
-    def flex_attn_with_mask(q, k, v, scale=None):
-        return flex_attention(q, k, v, block_mask=mask, scale=scale)
-    
-    # Compile the wrapper function
-    compiled_flex_attn = torch.compile(flex_attn_with_mask)
-    
-    return compiled_flex_attn
 
-def flex_sliding_tile_attention(q_all, k_all, v_all, strategy, tile_size, 
-                           image_size, text_length, scale=None):
-    device = q_all.device
-    
-    # Get the compiled flex attention function (cached if called with same parameters)
-    compiled_flex_attn = get_compiled_flex_attention(
-        strategy, 
-        tile_size,
-        image_size, 
-        text_length, 
-        device
-    )
-    
-    
-    # Apply the compiled flex attention
-    output = compiled_flex_attn(q_all, k_all, v_all, scale=scale)
-
-        
-    return output
 
 def attention(
     q,
@@ -139,41 +83,16 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=
     encoder_sequence_length = encoder_query.size(1)
 
     if mask_strategy[0] is not None:
-        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)
+        query = torch.cat([tile(query, nccl_info.sp_size), encoder_query], dim=1).transpose(1, 2).contiguous()
+        key = torch.cat([tile(key, nccl_info.sp_size), encoder_key], dim=1).transpose(1, 2).contiguous()
+        value = torch.cat([tile(value, nccl_info.sp_size), encoder_value], dim=1).transpose(1, 2).contiguous()
 
         head_num = query.size(1)
         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)]
 
-        if sliding_tile_attention is not None:
-            hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2)
-        else:
-            print("Sliding Tile Attention not available. Using Flex Sliding Tile Attention.")
-            hidden_states = torch.empty_like(query)
-            strategy_to_heads = {}
-            for head_index in range(head_num):
-                strategy = tuple(windows[head_index])  # Convert list to tuple for dict key
-                if strategy not in strategy_to_heads:
-                    strategy_to_heads[strategy] = []
-                strategy_to_heads[strategy].append(head_index)
-            for strategy, heads in strategy_to_heads.items():                
-                # Gather all heads with this strategy
-                query_heads = torch.cat([query[:, head_idx:head_idx + 1, :, :] for head_idx in heads], dim=1)
-                key_heads = torch.cat([key[:, head_idx:head_idx + 1, :, :] for head_idx in heads], dim=1)
-                value_heads = torch.cat([value[:, head_idx:head_idx + 1, :, :] for head_idx in heads], dim=1)
-                
-                # Process all heads with this strategy at once
-                # processed_heads = selected_attn_processor[processor_idx](query_heads, key_heads, value_heads)
-                processed_heads = flex_sliding_tile_attention(query_heads, key_heads, value_heads, strategy, (6, 8, 8), (30, 48, 80), text_length)
-                
-                # Distribute results back to the correct positions
-                for i, head_idx in enumerate(heads):
-                    hidden_states[:, head_idx:head_idx + 1, :, :] = processed_heads[:, i:i + 1, :, :]
-                
-                hidden_states = hidden_states.transpose(1, 2)
+        hidden_states = sliding_tile_attention(query, key, value, windows, text_length).transpose(1, 2).contiguous()
     else:
         query = torch.cat([query, encoder_query], dim=1)
         key = torch.cat([key, encoder_key], dim=1)
@@ -202,4 +121,4 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=
     b, s, a, d = attn.shape
     attn = attn.reshape(b, s, -1)
 
-    return attn
+    return attn

fastvideo/models/hunyuan/modules/models.py

@@ -523,6 +523,9 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
     ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
         if guidance is None:
             guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16)
+        if mask_strategy is None:
+            mask_strategy = [[None] * len(self.heads_num)
+                             for _ in range(len(self.double_blocks) + len(self.single_blocks))]
         img = x = hidden_states
         text_mask = encoder_attention_mask
         t = timestep

fastvideo/sample/sample_t2v_hunyuan.py

@@ -207,4 +207,4 @@ if __name__ == "__main__":
     # process for vae sequence parallel
     if args.vae_sp and not args.vae_tiling:
         raise ValueError("Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True.")
-    main(args)
+    main(args)

scripts/inference/inference_hunyuan.sh

@@ -37,4 +37,4 @@ torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     --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
+    --vae-sp