Initial changes: Refactor Attention

(cherry picked from commit 5a7237abfd)

unsloth/kernels/attention/__init__.py

@@ -0,0 +1,46 @@
+from .selector import (
+    AttentionBackendRegistry, 
+    select_attention_backend,
+    register_attention_backend,
+)
+from .layer import UnifiedAttention, create_attention_mechanism
+from .backends import (
+    AttentionBackend,
+    FlashAttentionBackend,
+    FlashAttentionSoftcapBackend,
+    XFormersBackend,
+    SDPABackend,
+    FlexAttentionBackend,
+    VanillaAttentionBackend,
+    VanillaSoftcappingAttentionBackend,
+)
+
+# Register the backends
+register_attention_backend("flash_attention", lambda: FlashAttentionBackend.is_available())
+register_attention_backend("flash_attention_softcap", lambda: FlashAttentionSoftcapBackend.is_available())
+register_attention_backend("xformers", lambda: XFormersBackend.is_available())
+register_attention_backend("sdpa", lambda: SDPABackend.is_available())
+register_attention_backend("flex_attention", lambda: FlexAttentionBackend.is_available())
+register_attention_backend("vanilla", lambda: True)  # Always available
+register_attention_backend("vanilla_softcap", lambda: True)  # Always available
+
+__all__ = [
+    # Selector functions
+    "AttentionBackendRegistry",
+    "select_attention_backend",
+    "register_attention_backend",
+    
+    # Unified attention interface
+    "UnifiedAttention",
+    "create_attention_mechanism",
+    
+    # Backend classes
+    "AttentionBackend",
+    "FlashAttentionBackend",
+    "FlashAttentionSoftcapBackend",
+    "XFormersBackend",
+    "SDPABackend",
+    "FlexAttentionBackend",
+    "VanillaAttentionBackend",
+    "VanillaSoftcappingAttentionBackend",
+]

unsloth/kernels/attention/backends/base.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple, Any, Dict, Union
+import torch
+
+class AttentionBackend(ABC):
+    """
+    Base class for attention backends.
+    """
+    @classmethod
+    def is_available(cls) -> bool:
+        return True
+    
+    @abstractmethod
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        causal_mask: Optional[Any] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        pass
+
+    @property
+    def name(self) -> str:
+        return self.__class__.__name__

unsloth/kernels/attention/backends/flash_attention.py

@@ -0,0 +1,147 @@
+import torch
+from typing import Optional, Any, Tuple, Dict
+from .base import AttentionBackend
+from ....models._utils import HAS_FLASH_ATTENTION, HAS_FLASH_ATTENTION_SOFTCAPPING
+
+if HAS_FLASH_ATTENTION:
+    from flash_attn import flash_attn_func
+
+class FlashAttentionBackend(AttentionBackend):
+    """
+    Flash Attention backend implementation.
+    Uses Flash Attention for efficient attention computation.
+    """
+    
+    def __init__(self, softmax_scale: float = None):
+        self.softmax_scale = softmax_scale
+    
+    @classmethod
+    def is_available(cls) -> bool:
+        return HAS_FLASH_ATTENTION
+    
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        causal_mask: Optional[Any] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        window_size: Tuple[int, int] = (-1, -1),
+        dropout_p: float = 0.0,
+        **kwargs
+    ) -> torch.Tensor:
+        if attention_mask is not None:
+            raise ValueError("FlashAttentionBackend does not support attention masks. Use a different backend.")
+            
+        output = flash_attn_func(
+            query, key, value,
+            causal=True,
+            window_size=window_size,
+            softmax_scale=self.softmax_scale,
+            dropout_p=dropout_p
+        )
+        
+        batch_size, seq_len, num_heads, head_dim = output.shape
+        output = output.reshape(batch_size, seq_len, num_heads * head_dim)
+        
+        return output
+    
+    @classmethod
+    def supports_feature(cls, feature: str) -> bool:
+        """
+        Check if this backend supports a particular feature.
+        """
+        features = {
+            "causal": True,
+            "sliding_window": True,
+            "attention_mask": False,
+            "dropout": True,
+            "custom_scale": True,
+        }
+        return features.get(feature, False)
+
+
+class FlashAttentionSoftcapBackend(FlashAttentionBackend):
+    """
+    Flash Attention backend with softcapping support.
+    Used primarily for Gemma 2 models.
+    """
+    
+    def __init__(self, softmax_scale: float = None, softcap: float = None):
+        """
+        Initialize the Flash Attention Softcap backend.
+        
+        Args:
+            softmax_scale: Optional scaling factor for the softmax operation
+            softcap: Softcap value for attention logits
+        """
+        super().__init__(softmax_scale)
+        self.softcap = softcap
+    
+    @classmethod
+    def is_available(cls) -> bool:
+        return HAS_FLASH_ATTENTION_SOFTCAPPING
+    
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        causal_mask: Optional[Any] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        window_size: Tuple[int, int] = (-1, -1),
+        dropout_p: float = 0.0,
+        softcap: float = None,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Perform attention using Flash Attention with softcapping.
+        
+        Args:
+            query: Query tensor of shape [batch_size, num_heads, seq_len, head_dim]
+            key: Key tensor of shape [batch_size, num_kv_heads, seq_len, head_dim]
+            value: Value tensor of shape [batch_size, num_kv_heads, seq_len, head_dim]
+            causal_mask: Optional causal mask
+            attention_mask: Optional attention mask
+            window_size: Sliding window size as (backward, forward) tuple
+            dropout_p: Dropout probability
+            softcap: Softcap value for attention logits
+            **kwargs: Additional arguments
+            
+        Returns:
+            torch.Tensor: Output tensor after attention
+        """
+        if attention_mask is not None:
+            raise ValueError("FlashAttentionSoftcapBackend does not support attention masks. Use a different backend.")
+            
+        softcap_value = softcap if softcap is not None else self.softcap
+        
+        output = flash_attn_func(
+            query, key, value,
+            causal=True,
+            softcap=softcap_value,
+            softmax_scale=self.softmax_scale,
+            window_size=window_size,
+            dropout_p=dropout_p
+        )
+        
+        batch_size, seq_len, num_heads, head_dim = output.shape
+        output = output.reshape(batch_size, seq_len, num_heads * head_dim)
+        
+        return output
+    
+    @classmethod
+    def supports_feature(cls, feature: str) -> bool:
+        """
+        Check if this backend supports a particular feature.
+        
+        Args:
+            feature: Name of the feature to check
+            
+        Returns:
+            bool: True if the feature is supported, False otherwise
+        """
+        features = super().supports_feature(feature)
+        if feature == "softcap":
+            return True
+        return features

unsloth/kernels/attention/layer.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple, Dict, Any, Union
+import logging
+
+from .selector import select_attention_backend
+from .backends import AttentionBackend
+
+logger = logging.getLogger(__name__)
+
+class UnifiedAttention(nn.Module):
+    """
+    Unified attention layer that delegates to the appropriate attention backend.
+    This provides a consistent interface for all attention mechanisms in the codebase.
+    """
+    
+    def __init__(
+        self,
+        config: Optional[Any] = None, 
+        backend: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Initialize the unified attention layer.
+        
+        Args:
+            config: Model configuration object
+            backend: Name of the specific backend to use (if None, auto-select)
+            **kwargs: Additional arguments to pass to the backend
+        """
+        super().__init__()
+        
+        backend_name, backend_cls = select_attention_backend(backend, config, **kwargs)
+        logger.debug(f"Selected attention backend: {backend_name}")
+        
+        self.backend = backend_cls(**kwargs)
+        self.backend_name = backend_name
+        
+        # Store config for potential future use
+        self.config = config
+    
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        causal_mask: Optional[Any] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Perform attention using the selected backend.
+        
+        Args:
+            query: Query tensor [batch_size, num_heads, seq_len, head_dim]
+            key: Key tensor [batch_size, num_kv_heads, seq_len, head_dim]
+            value: Value tensor [batch_size, num_kv_heads, seq_len, head_dim]
+            causal_mask: Optional causal mask
+            attention_mask: Optional attention mask
+            **kwargs: Additional arguments to pass to the backend
+            
+        Returns:
+            torch.Tensor: Output tensor after attention
+        """
+        return self.backend.forward(
+            query=query,
+            key=key,
+            value=value,
+            causal_mask=causal_mask,
+            attention_mask=attention_mask,
+            **kwargs
+        )
+    
+    @property
+    def name(self) -> str:
+        """Get the name of the currently used backend."""
+        return self.backend_name
+    
+    def supports_feature(self, feature: str) -> bool:
+        """Check if the current backend supports a specific feature."""
+        return self.backend.supports_feature(feature)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the attention configuration to a dictionary."""
+        return {
+            "backend": self.backend_name,
+            "config": self.config.__dict__ if hasattr(self.config, "__dict__") else None
+        }
+
+def create_attention_mechanism(
+    config: Optional[Any] = None,
+    backend: Optional[str] = None,
+    **kwargs
+) -> UnifiedAttention:
+    """
+    Factory function to create an attention mechanism.
+    
+    Args:
+        config: Model configuration object
+        backend: Name of the specific backend to use (if None, auto-select)
+        **kwargs: Additional arguments to pass to the backend
+        
+    Returns:
+        UnifiedAttention: An initialized attention layer
+    """
+    return UnifiedAttention(config, backend, **kwargs)

unsloth/kernels/attention/selector.py

@@ -0,0 +1,64 @@
+import os
+import logging
+from functools import cache
+from typing import Dict, Callable, Generator, Optional, Type, Any, Union, Tuple
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Registry to hold all attention backends
+AttentionBackendRegistry = {}
+
+def register_attention_backend(name: str, condition_fn: Callable[[], bool] = None):
+    """
+    Decorator to register an attention backend
+    """
+    def decorator(cls):
+        if condition_fn is None or condition_fn():
+            AttentionBackendRegistry[name] = cls
+            logger.debug(f"Registeres attention backedn: {name}")
+        return cls
+    return decorator
+
+def select_attention_backend(
+        backend_name: Optional[str] = None,
+        config: Optional[Any] = None,
+        **kwargs
+) -> Tuple[str, Any]:
+    """
+    Select the appropriate attention backend based on the provided criteria
+    """
+
+    if backend_name is not None:
+        if backend_name in AttentionBackendRegistry:
+            return backend_name, AttentionBackendRegistry[backend_name]
+        else:
+            logger.warning(f"Given backend not available: {backend_name}")
+
+    available_backends = list(AttentionBackendRegistry.keys())
+
+    priority_order = os.environ.get("GLOBAL_ATTENTION_PRIORITY", "").split(",")
+    if not priority_order or priority_order == [""]:
+        priority_order = ["flash_attention", "flex_attention", "xformers", "sdpa"]
+    
+    priority_order.extend(b for b in available_backends if b not in priority_order)
+
+    if config is not None:
+        # Softcapping need for Gemma 2
+        has_softcapping = hasattr(config, "attn_logit_softcapping") and config.attn_logit_softcapping > 0
+        if has_softcapping and "flash_attention_softcapping" in available_backends:
+            return "flash_attention_softcapping", AttentionBackendRegistry["flash_attention_softcap"]
+        
+        # Sliding Window Attention
+        has_swa = hasattr(config, "sliding_window") and config.sliding_window not in (None, "null")
+        if has_swa:
+            for backend in ["flash_attention", "flex_attention"]:
+                if backend in available_backends:
+                    return backend, AttentionBackendRegistry[backend]
+                
+    for backend in priority_order:
+        if backend in available_backends:
+            return backend, AttentionBackendRegistry[backend]
+        
+    raise RuntimeError("No attention backends available.")