Memory-efficient attention using linear attention mechanism. Supports automatic fallback to optimized implementations when available.
| 616 | return should_use_mlx(model_id) |
| 617 | |
| 618 | class MemoryEfficientAttention(nn.Module): |
| 619 | """ |
| 620 | Memory-efficient attention using linear attention mechanism. |
| 621 | Supports automatic fallback to optimized implementations when available. |
| 622 | """ |
| 623 | def __init__( |
| 624 | self, |
| 625 | hidden_size: int, |
| 626 | num_attention_heads: int, |
| 627 | dropout: float = 0.1, |
| 628 | ): |
| 629 | super().__init__() |
| 630 | self.hidden_size = hidden_size |
| 631 | self.num_attention_heads = num_attention_heads |
| 632 | self.head_dim = hidden_size // num_attention_heads |
| 633 | self.scale = self.head_dim ** -0.5 |
| 634 | |
| 635 | # Standard projections with bias |
| 636 | self.q_proj = nn.Linear(hidden_size, hidden_size) |
| 637 | self.k_proj = nn.Linear(hidden_size, hidden_size) |
| 638 | self.v_proj = nn.Linear(hidden_size, hidden_size) |
| 639 | self.o_proj = nn.Linear(hidden_size, hidden_size) |
| 640 | |
| 641 | self.dropout = nn.Dropout(dropout) |
| 642 | |
| 643 | # Try to import optimized attention implementations |
| 644 | self.optimized_attention = None |
| 645 | |
| 646 | # Try Flash Attention |
| 647 | try: |
| 648 | from flash_attn import flash_attn_func |
| 649 | self.optimized_attention = flash_attn_func |
| 650 | print("Using Flash Attention") |
| 651 | except ImportError: |
| 652 | pass |
| 653 | |
| 654 | # Try xFormers |
| 655 | if self.optimized_attention is None: |
| 656 | try: |
| 657 | import xformers.ops as xops |
| 658 | self.optimized_attention = xops.memory_efficient_attention |
| 659 | print("Using xFormers attention") |
| 660 | except ImportError: |
| 661 | pass |
| 662 | |
| 663 | def _linear_attention( |
| 664 | self, |
| 665 | q: torch.Tensor, |
| 666 | k: torch.Tensor, |
| 667 | v: torch.Tensor, |
| 668 | attention_mask: Optional[torch.Tensor] = None, |
| 669 | causal: bool = False, |
| 670 | ) -> torch.Tensor: |
| 671 | """Implements linear attention for more memory efficiency""" |
| 672 | # Scale query |
| 673 | q = q * self.scale |
| 674 | |
| 675 | # Apply mask if provided |
no outgoing calls
no test coverage detected