Fix sageattn_varlen

wanvideo/modules/attention.py

@@ -14,11 +14,11 @@ try:
     FLASH_ATTN_2_AVAILABLE = True
 except Exception as e:
     FLASH_ATTN_2_AVAILABLE = False
-        
+
 # Sage Attention imports
 try:
     from sageattention import sageattn
-    
+
     @torch.library.custom_op("wanvideo::sageattn", mutates_args=())
     def sageattn_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, tensor_layout: str = "HND"
     ) -> torch.Tensor:
@@ -28,7 +28,7 @@ try:
             return sageattn(q.to(torch.float16), k.to(torch.float16), v.to(torch.float16), attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, tensor_layout=tensor_layout).to(torch.float32)
         else:
             return sageattn(q, k, v, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, tensor_layout=tensor_layout)
-    
+
     @sageattn_func.register_fake
     def _(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False, tensor_layout="HND"):
         # Return tensor with same shape as q
@@ -49,12 +49,12 @@ except Exception as e:
         log.warning("sageattention DLL loading error, sageattention will not be available")
     sageattn_func = None
 
-try: 
+try:
     from sageattention import sageattn_varlen
-    
+    from typing import List
+
     @torch.library.custom_op("wanvideo::sageattn_varlen", mutates_args=())
-    def sageattn_varlen_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_lens: list, k_lens: list, max_seqlen_q: int, max_seqlen_k: int, dropout_p: float = 0.0, is_causal: bool = False
-    ) -> torch.Tensor:
+    def sageattn_varlen_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_lens: List[int], k_lens: List[int], max_seqlen_q: int, max_seqlen_k: int, dropout_p: float = 0.0, is_causal: bool = False) -> torch.Tensor:
         cu_seqlens_q = torch.tensor([0] + list(torch.cumsum(torch.tensor(q_lens), dim=0)), device=q.device, dtype=torch.int32)
         cu_seqlens_k = torch.tensor([0] + list(torch.cumsum(torch.tensor(k_lens), dim=0)), device=q.device, dtype=torch.int32)
         if not (q.dtype == k.dtype == v.dtype):
@@ -63,12 +63,12 @@ try:
             return sageattn_varlen(q.to(torch.float16), k.to(torch.float16), v.to(torch.float16), cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p=dropout_p, is_causal=is_causal).to(torch.float32)
         else:
             return sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p=dropout_p, is_causal=is_causal)
-    
+
     @sageattn_varlen_func.register_fake
     def _(q, k, v, q_lens, k_lens, max_seqlen_q, max_seqlen_k, dropout_p=0.0, is_causal=False):
         # Return tensor with same shape as q
         return q.clone()
-except: 
+except Exception as e:
     sageattn_varlen_func = None
 
 try:
@@ -210,7 +210,7 @@ def attention(
     dtype=torch.bfloat16,
     attention_mode='sdpa',
     attn_mask=None,
-):  
+):
     if "flash" in attention_mode:
         if attention_mode == 'flash_attn_2':
             fa_version = 2
@@ -237,9 +237,9 @@ def attention(
         return torch.nn.functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=attn_mask).transpose(1, 2).contiguous()
     elif attention_mode == 'sageattn_3':
         return sageattn_blackwell(
-            q.transpose(1,2), 
-            k.transpose(1,2), 
-            v.transpose(1,2), 
+            q.transpose(1,2),
+            k.transpose(1,2),
+            v.transpose(1,2),
             per_block_mean=False #seems necessary for reasonable VRAM usage, not sure of other implications
             ).transpose(1,2).contiguous()
     elif attention_mode == 'sageattn_varlen':