From b8eee7a8ba6f4b52fbc7be6d870b676ebe0d0c07 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 14 Mar 2026 05:31:14 +0000 Subject: [PATCH] Revert "Initial changes: Refactor Attention" This reverts commit a2af843271169c0efafdaaf4a172971f65d6958e. --- unsloth/kernels/attention/__init__.py | 46 ------ unsloth/kernels/attention/backends/base.py | 27 ---- .../attention/backends/flash_attention.py | 147 ------------------ unsloth/kernels/attention/backends/sdpa.py | 0 .../kernels/attention/backends/xformers.py | 0 unsloth/kernels/attention/layer.py | 106 ------------- unsloth/kernels/attention/selector.py | 64 -------- .../backends => }/flex_attention.py | 0 8 files changed, 390 deletions(-) delete mode 100644 unsloth/kernels/attention/__init__.py delete mode 100644 unsloth/kernels/attention/backends/base.py delete mode 100644 unsloth/kernels/attention/backends/flash_attention.py delete mode 100644 unsloth/kernels/attention/backends/sdpa.py delete mode 100644 unsloth/kernels/attention/backends/xformers.py delete mode 100644 unsloth/kernels/attention/layer.py delete mode 100644 unsloth/kernels/attention/selector.py rename unsloth/kernels/{attention/backends => }/flex_attention.py (100%) diff --git a/unsloth/kernels/attention/__init__.py b/unsloth/kernels/attention/__init__.py deleted file mode 100644 index 996de41566..0000000000 --- a/unsloth/kernels/attention/__init__.py +++ /dev/null @@ -1,46 +0,0 @@ -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", -] \ No newline at end of file diff --git a/unsloth/kernels/attention/backends/base.py b/unsloth/kernels/attention/backends/base.py deleted file mode 100644 index 63b6bb2a99..0000000000 --- a/unsloth/kernels/attention/backends/base.py +++ /dev/null @@ -1,27 +0,0 @@ -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__ \ No newline at end of file diff --git a/unsloth/kernels/attention/backends/flash_attention.py b/unsloth/kernels/attention/backends/flash_attention.py deleted file mode 100644 index 87911c51dd..0000000000 --- a/unsloth/kernels/attention/backends/flash_attention.py +++ /dev/null @@ -1,147 +0,0 @@ -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 \ No newline at end of file diff --git a/unsloth/kernels/attention/backends/sdpa.py b/unsloth/kernels/attention/backends/sdpa.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/unsloth/kernels/attention/backends/xformers.py b/unsloth/kernels/attention/backends/xformers.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/unsloth/kernels/attention/layer.py b/unsloth/kernels/attention/layer.py deleted file mode 100644 index a21b12e849..0000000000 --- a/unsloth/kernels/attention/layer.py +++ /dev/null @@ -1,106 +0,0 @@ -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) diff --git a/unsloth/kernels/attention/selector.py b/unsloth/kernels/attention/selector.py deleted file mode 100644 index 03362b34a5..0000000000 --- a/unsloth/kernels/attention/selector.py +++ /dev/null @@ -1,64 +0,0 @@ -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.") diff --git a/unsloth/kernels/attention/backends/flex_attention.py b/unsloth/kernels/flex_attention.py similarity index 100% rename from unsloth/kernels/attention/backends/flex_attention.py rename to unsloth/kernels/flex_attention.py