Refactor fdma_peft_integration_check and _flash_dynamic_mask_attention_forward for clarity and consistency; rename keep_window_size to window_size and enhance FlashDynamicMaskAttentionKwargs documentation.

Author: algo-home

flash_dmattn/integrations/modeling_flash_dynamic_mask_attention_utils.py

@@ -17,16 +17,29 @@
 from .import_utils import is_flash_dmattn_available
 
 from transformers.utils import logging
-from transformers.integrations import flash_attention
 
 
 logger = logging.get_logger(__name__)
 
 
-def fdma_peft_integration_check(q, k, v, bias, target_dtype: Optional[torch.dtype] = None):
+def fdma_peft_integration_check(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    target_dtype: Optional[torch.dtype] = None
+):
+    """
+    PEFT usually casts the layer norms in float32 for training stability reasons
+    therefore the input hidden states gets silently casted in float32. Hence, we need
+    cast them back in float16 / bfloat16 just to be sure everything works as expected.
+    This might slowdown training & inference so it is recommended to not cast the LayerNorms!
+    """
     if target_dtype and q.dtype == torch.float32:
         logger.warning_once(f"Casting fp32 inputs back to {target_dtype} for flash-dmattn compatibility.")
-        q, k, v, bias = q.to(target_dtype), k.to(target_dtype), v.to(target_dtype), bias.to(target_dtype)
+        q, k, v = q.to(target_dtype), k.to(target_dtype), v.to(target_dtype)
+        if bias is not None:
+            bias = bias.to(target_dtype)
     return q, k, v, bias
 
 
@@ -43,8 +56,24 @@ def _lazy_imports(impl: Optional[str]):
 
 
 class FlashDynamicMaskAttentionKwargs(TypedDict, total=False):
-    cumulative_seqlens_q: Optional[torch.LongTensor]
-    cumulative_seqlens_k: Optional[torch.LongTensor]
+    """
+    Keyword arguments for Flash Dynamic Mask Attention with Compile.
+
+    Attributes:
+        cu_seq_lens_q (`torch.LongTensor`, *optional*)
+            Gets cumulative sequence length for query state.
+        cu_seq_lens_k (`torch.LongTensor`, *optional*)
+            Gets cumulative sequence length for key state.
+        max_length_q (`int`, *optional*):
+            Maximum sequence length for query state.
+        max_length_k (`int`, *optional*):
+            Maximum sequence length for key state.
+    """
+
+    cu_seq_lens_q: Optional[torch.LongTensor]
+    cu_seq_lens_k: Optional[torch.LongTensor]
+    max_length_q: Optional[int]
+    max_length_k: Optional[int]
 
 
 def _flash_dynamic_mask_attention_forward(
@@ -58,15 +87,14 @@ def _flash_dynamic_mask_attention_forward(
     is_causal: bool,
     softmax_scale: Optional[float] = None,
     softcap: Optional[float] = None,
-    keep_window_size: Optional[int] = None,
+    window_size: Optional[int] = None,
     deterministic: Optional[bool] = None,
     target_dtype: Optional[torch.dtype] = None,
     implementation: Optional[str] = None,
     **kwargs,
 ):
     dtype = query_states.dtype
     min_dtype = torch.finfo(dtype).min
-    batch_size, _, num_kv_heads, _ = key_states.shape
 
     if not all(k in globals() for k in ("_flash_fn")):
         flash_fn = _lazy_imports(implementation)
@@ -93,14 +121,12 @@ def _flash_dynamic_mask_attention_forward(
             min_dtype
         )
 
-    if keep_window_size is not None and key_length > keep_window_size:
+    if window_size is not None and key_length > window_size:
         topk_values, topk_indices = torch.topk(
-            attention_bias, keep_window_size, dim=-1, largest=True, sorted=False
+            attention_bias, window_size, dim=-1, largest=True, sorted=False
         )
         attention_mask = torch.zeros_like(attention_bias, dtype=torch.bool, device=attention_bias.device)
         attention_mask = attention_mask.scatter(-1, topk_indices, topk_values != min_dtype)
-    else:
-        attention_mask = None
 
     out = flash_fn(
         query_states, key_states, value_states, attn_mask=attention_mask, attn_bias=attention_bias, scale=softmax_scale, is_causal=is_causal