Add support for FlashAttention

README.md

@@ -54,6 +54,7 @@ This model is intended for:
 ### Image Classification Example
 
 ```python
+import torch
 from transformers import AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 import requests
@@ -72,12 +73,14 @@ model = AutoModelForImageClassification.from_pretrained(
     "XAFT/SM-Selective-ViT-Tiny-Tall-224",
     trust_remote_code=True,
 )
+model = model.half() # Cast to FP16 to enable FlashAttention
 
 # Preprocess
 inputs = processor(
     images=image,
     return_tensors="pt",
 )
+inputs = inputs.to(torch.half) # Cast to FP16
 
 # Forward pass
 outputs = model(**inputs)

modeling_selectivevit.py

@@ -50,6 +50,7 @@ class SMSelectiveViTModelForClassification(PreTrainedModel ):
         full=False,
         output_hidden_states=None,
         return_dict=None,
+        skip_masks=False,
         **kwargs,
     ):
         output_hidden_states = (
@@ -67,6 +68,7 @@ class SMSelectiveViTModelForClassification(PreTrainedModel ):
             pixel_values,
             full=full,
             output_hidden_states=output_hidden_states,
+            skip_masks=skip_masks
         )
 
         logits, distil_logits = self.backbone.forward_classifier(last_hidden)

selective_vit.py

@@ -1,8 +1,23 @@
+import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchvision.ops import StochasticDepth
 import math
+import warnings
+try:
+    import torch.nn.attention.varlen as varlen
+    HAS_VARLEN_FLASH_ATTENTION = True
+except ImportError:
+    warnings.warn(
+        "Could not import torch.nn.attention.varlen, variable length Flash Attention is disabled.",
+        category=UserWarning,
+        stacklevel=2,
+    )
+    HAS_VARLEN_FLASH_ATTENTION = False
+
+enable_fa = os.environ.get('DISABLE_FA', '0').lower() not in {"1", "true", "yes", "y", "on"}
+HAS_VARLEN_FLASH_ATTENTION = HAS_VARLEN_FLASH_ATTENTION and enable_fa
 
 class SoftMaskedMultiheadAttention(nn.Module):
     def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True,
@@ -47,12 +62,9 @@ class SoftMaskedMultiheadAttention(nn.Module):
         if self.v_proj.bias is not None:
             nn.init.constant_(self.out_proj.bias, 0.)
 
-    def forward(self, query, key, value, key_padding_mask=None,
-                need_weights=True, attn_mask=None, average_attn_weights=True):
-        """
-        query, key, value: shape (L, N, E)
-        where L is the sequence length, N is the batch size, E is the embedding dimension.
-        """
+
+    def naive_forward(self, query, key, value, key_padding_mask=None,
+            attn_mask=None, average_attn_weights=True):
         batch_size, tgt_len, embed_dim = query.size()
         batch_size, src_len, _ = key.size()
 
@@ -74,10 +86,11 @@ class SoftMaskedMultiheadAttention(nn.Module):
             # Ensure attn_mask values are in (0, 1] to avoid log(0)
             # attn_mask shape [b, l]
             attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
-            if not self.training:
-                scores = scores.masked_fill((attn_mask == 0.), float('-inf'))
             eps = 1e-6
-            attn_mask = attn_mask.clip(min=eps).log()
+            attn_mask_l = attn_mask.clip(min=eps).log()
+            if not self.training:
+                attn_mask_l = attn_mask_l.masked_fill((attn_mask == 0.), float('-inf'))
+            attn_mask = attn_mask_l
             # attn_mask shape [b, 1, 1, l]
             scores = scores + self.scale * attn_mask
 
@@ -99,16 +112,96 @@ class SoftMaskedMultiheadAttention(nn.Module):
 
         attn_output = self.out_proj(attn_output)
 
-        if need_weights:
-            # Optionally average attention weights over heads
-            if average_attn_weights:
-                attn_weights = attn_weights.mean(dim=1)
-            else:
-                attn_weights = attn_weights
+        return attn_output
+
+    def flash_forward(
+        self,
+        query, key, value,
+        cu_seq_q, cu_seq_k,
+        max_q, max_k,
+        attn_mask=None,
+        is_causal=False,
+    ):
+        """
+        FlashAttention-compatible soft-masked attention using varlen_attn
+        """
+
+        q = self.q_proj(query)   # (Tq, H*D)
+        k = self.k_proj(key)     # (Tk, H*D)
+        v = self.v_proj(value)   # (Tk, H*D)
+
+        Tq = q.shape[0]
+        Tk = k.shape[0]
+
+        q = q.view(Tq, self.num_heads, self.head_dim)
+        k = k.view(Tk, self.num_heads, self.head_dim)
+        v = v.view(Tk, self.num_heads, self.head_dim)
+
+        # Apply the soft [0, 1] mask
+        if attn_mask is not None:
+            # attn_mask: (Tk,) or (B, Lk) flattened to match Tk
+            eps = 1e-6
+            attn_mask_l = attn_mask.clip(min=eps).log()
+            if not self.training: # Inference mode can have infinite atten scores
+                attn_mask_l = attn_mask_l.masked_fill((attn_mask == 0.), float('-inf'))
+            log_m = attn_mask_l
+
+            # Broadcast to (Tk, H, 1)
+            log_m = log_m.view(Tk, 1, 1).expand(-1, self.num_heads, 1)
+            k_zeros = torch.zeros_like(log_m).expand(-1, -1, 7)
+
+            # Augment K and Q
+            # We want:
+            #   (qk^T)/sqrt(d) + scale * log(m)
+            scale_attn = 1.0 / math.sqrt(self.head_dim)
+
+            k_extra = log_m * (self.scale / scale_attn)
+
+            k = torch.cat([k, k_extra, k_zeros], dim=-1)     # (Tk, H, D+1)
+
+            v_zeros = torch.zeros(Tk, self.num_heads, 8, device=v.device, dtype=v.dtype)
+            v = torch.cat([v, v_zeros], dim=-1)
+
+            q_ones = torch.ones(
+                Tq, self.num_heads, 8,
+                device=q.device, dtype=q.dtype
+            )
+            q = torch.cat([q, q_ones], dim=-1)      # (Tq, H, D+1)
+
+            attn_dim = self.head_dim + 1
+        else:
+            attn_dim = self.head_dim
+            scale_attn = 1.0 / math.sqrt(self.head_dim)
+
+        # FlashAttention varlen call
+        out = varlen.varlen_attn(
+            query=q,
+            key=k,
+            value=v,
+            cu_seq_q=cu_seq_q,
+            cu_seq_k=cu_seq_k,
+            max_q=max_q,
+            max_k=max_k,
+            is_causal=is_causal,
+            scale=scale_attn,
+        )
+
+        # Merge heads and output projection
+        out = out[..., :self.head_dim]
+        out = out.reshape(Tq, self.num_heads * self.head_dim)
+        out = self.out_proj(out)
+
+        return out
+
+    def forward(self, query, key, value, method="naive", **kwargs):
+        if method == 'naive':
+            out = self.naive_forward(query, key, value, **kwargs)
+        elif method == "fa":
+            out = self.flash_forward(query, key, value, **kwargs)
         else:
-            attn_weights = None
+            raise ValueError(f"No attention method named {method}.")
+        return out
 
-        return attn_output, attn_weights
 
 def get_ffn(input_dim, output_dim, middle_dim, dropout=0.1):
     fc1 = nn.Linear(input_dim, middle_dim)
@@ -122,6 +215,7 @@ def get_ffn(input_dim, output_dim, middle_dim, dropout=0.1):
             nn.Dropout(dropout),
             fc3
             )
+
 # Assuming SoftMaskedMultiheadAttention is already defined as provided earlier
 class EncoderBlock(nn.Module):
     def __init__(self, input_dim, embed_dim, num_heads, mlp_dim, dropout=0.1, drop_path=0.0, patch_drop=0.0, attention_scale=2., mask_threshold=0.05):
@@ -162,16 +256,13 @@ class EncoderBlock(nn.Module):
         nn.init.zeros_(self.norm3.weight)
         nn.init.zeros_(self.norm3.bias)
 
-    def forward_common(self, x, mask):
-        """
-        x: shape (batch_size, seq_len, embed_dim)
-        """
+    def forward_common(self, x, mask, skip_masks=False):
         # Compute mask scores: (batch_size, seq_len, 1)
         x1 = x
         x = self.embed(x)
         x = self.norm1(x)
         # Apply attention mechanism
-        attn_output, _ = self.self_attn(x, x, x, attn_mask=mask)
+        attn_output = self.self_attn(x, x, x, attn_mask=mask if not skip_masks else None, method="naive")
         # Add & Norm
         x = x + self.path_drop(attn_output)
         x = self.norm2(x)
@@ -185,6 +276,45 @@ class EncoderBlock(nn.Module):
         x = x1 + x
         return x
 
+    def flash_forward(self, x, mask, skip_masks=False):
+        binary_mask = mask >= self.mask_threshold
+        sel_mask = mask[binary_mask]
+
+        seq_lengths = binary_mask.sum(1)
+        cum_lengths = torch.zeros(binary_mask.shape[0]+1, dtype=torch.int, device=binary_mask.device)
+        cum_lengths[1:] = seq_lengths.cumsum(-1)
+        max_len = seq_lengths.amax()
+
+        x1 = x
+        x = x[binary_mask]
+
+        x = self.embed(x)
+        x = self.norm1(x)
+        # Apply flash attention mechanism
+        attn_output = self.self_attn(
+                x, x, x,
+                cu_seq_q=cum_lengths,
+                cu_seq_k=cum_lengths,
+                max_q=max_len,
+                max_k=max_len,
+                attn_mask=sel_mask if not skip_masks else None,
+                method="fa"
+                )
+        # Add & Norm
+        x = x + self.path_drop(attn_output)
+        x = self.norm2(x)
+        # Feed-forward network
+        mlp_output = self.mlp(x)
+        # Add & Norm
+        x = self.path_drop(self.project(x + mlp_output))
+        x = self.norm3(x)
+        if mask is not None:
+            x = x * sel_mask.unsqueeze(-1)
+
+        x_out = x1.clone()
+        x_out[binary_mask] = x_out[binary_mask] + x
+        return x_out
+
     def get_groups(self, mask, full=False):
         n_items, index = (mask != 0.0).sum(-1).cpu().sort(descending=True)
         n_items, index = n_items.tolist(), index.tolist()
@@ -198,13 +328,7 @@ class EncoderBlock(nn.Module):
             groups[-1][0].append(ii)
         return groups
 
-    def infer_forward(self, x, mask, full=False):
-        """
-        The “sparse‐inference” path: for each group of batch‐samples that have the same
-        number n of tokens ≥ mask_threshold, gather only those top‐n tokens (in original order),
-        run forward_common on the smaller (b’, n, dim) tensor, then scatter the results back.
-        Fully masked tokens are left untouched.
-        """
+    def naive_forward(self, x, mask, full=False, skip_masks=False):
         # Step 1: Threshold the mask without in-place ops
         mask_thresholded = mask * (mask >= self.mask_threshold)
         # Step 2: Prepare output tensor (copy of x)
@@ -223,7 +347,7 @@ class EncoderBlock(nn.Module):
             X_topk = torch.gather(x_sel, dim=1, index=idx_expanded)
             mask_topk = torch.gather(mask_sel, dim=1, index=topk_idx_sorted)
             # Run forward pass
-            results = self.forward_common(X_topk, mask_topk)
+            results = self.forward_common(X_topk, mask_topk, skip_masks)
             # Scatter results into a new x_sel tensor
             x_sel_updated = x_sel.clone()
             x_sel_updated = x_sel_updated.scatter(1, idx_expanded, results)
@@ -231,15 +355,27 @@ class EncoderBlock(nn.Module):
             x_out[batch_indices] = x_sel_updated
         return x_out
 
-    def forward(self, x, full=False):
+    def forward(self, x, full=False, skip_masks=False):
         if self.linear_mask is not None:
             attn_mask = self.patch_drop(self.linear_mask(x).sigmoid().squeeze(-1))
         else:
             attn_mask = None
         if not self.training and not attn_mask is None and self.mask_threshold >= 0:
-            x = self.infer_forward(x, attn_mask, full)
+            if (
+                    HAS_VARLEN_FLASH_ATTENTION and
+                    'cuda' in x.device.type and
+                    x.dtype in (torch.bfloat16, torch.float16)
+                    ):
+                x = self.flash_forward(x, attn_mask, skip_masks)
+            else:
+                warnings.warn(
+                    "Flash Attention requirements not met, falling back to naive attention.",
+                    category=UserWarning,
+                    stacklevel=2,
+                )
+                x = self.naive_forward(x, attn_mask, full, skip_masks)
         else:
-            x = self.forward_common(x, attn_mask)
+            x = self.forward_common(x, attn_mask, skip_masks)
         return x, attn_mask
 
 
@@ -329,15 +465,8 @@ class VisionTransformer(nn.Module):
         pixel_values,
         full=False,
         output_hidden_states=False,
+        skip_masks=False
     ):
-        """
-        Args:
-            pixel_values: (B, C, H, W)
-        Returns:
-            last_hidden_state: (B, N, D)
-            all_hidden_states: tuple or None
-            masks: Tensor or None
-        """
         batch_size = pixel_values.size(0)
         hidden_states = []
 
@@ -361,7 +490,7 @@ class VisionTransformer(nn.Module):
         masks = []
 
         for layer in self.encoder_layers:
-            x, mask = layer(x, full)
+            x, mask = layer(x, full, skip_masks=skip_masks)
 
             if output_hidden_states:
                 hidden_states.append(x)
@@ -385,13 +514,6 @@ class VisionTransformer(nn.Module):
 
 
     def forward_classifier(self, hidden_states):
-        """
-        Args:
-            hidden_states: (B, N, D)
-        Returns:
-            logits: (B, num_classes)
-            dis_logits: (B, num_classes) or None
-        """
         cls_token = hidden_states[:, 0]
         logits = self.head(cls_token)
 
@@ -406,7 +528,7 @@ class VisionTransformer(nn.Module):
 
         return logits, dis_logits
 
-    def forward(self, x, full=False):
-        last_hidden_states, hidden_states, masks = self.forward_features(x, full)
+    def forward(self, x, full=False, skip_masks=False):
+        last_hidden_states, hidden_states, masks = self.forward_features(x, full, skip_masks=skip_masks)
         logits, dis_logits = self.forward_classifier(last_hidden_states)
         return logits, dis_logits, masks