Optimizes attention mask handling with boolean dtype

Replaces float-based attention mask operations with boolean dtype for improved memory efficiency and cleaner logic.

Removes unnecessary dtype conversion and simplifies mask creation by using boolean tensors directly instead of converting comparison results to float values.

Author: algo-home

flash_dmattn/integrations/modeling_flash_dynamic_mask_attention_utils.py

@@ -93,17 +93,15 @@ def _flash_dynamic_mask_attention_forward(
             ~attention_mask,
             min_dtype
         )
-        attention_mask = attention_mask.to(dtype)
 
     if keep_window_size is not None:
         if key_length > keep_window_size:
             topk_values, topk_indices = torch.topk(
                 attention_bias, keep_window_size, dim=-1, largest=True, sorted=False
             )
-            valid_topk = (topk_values != min_dtype).to(dtype)
-            attention_mask = torch.zeros_like(attention_bias, dtype=dtype, device=attention_bias.device)
-            attention_mask = attention_mask.scatter(-1, topk_indices, valid_topk)
-            attention_bias = attention_bias.masked_fill(attention_mask == 0.0, min_dtype)
+            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)
+            attention_bias = attention_bias.masked_fill(attention_mask == False, min_dtype)
 
     out = flash_fn(
         query_states, key_states, value_states, attn_mask=attention_mask, attn_bias=attention_bias, scale=softmax_scale, is_causal=is_causal