File size: 3,545 Bytes
b2140f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """
Mesosfer Bear AI - Flash Attention Backend
Provides a unified attention function that auto-selects the best backend:
1. flash-attn package (Dao-AILab / ROCm Triton) — fastest, native GQA
2. PyTorch SDPA — auto-dispatches to native FlashAttention / AOTriton kernels
Usage:
from engine.flashattion import bear_attention, get_attn_backend
"""
import os
from typing import Optional
import torch
import torch.nn.functional as F
def _is_rocm() -> bool:
"""Return True if running on AMD ROCm (HIP)."""
return bool(getattr(torch.version, "hip", None))
# -- Backend detection ------------------------------------------------------
if _is_rocm():
os.environ.setdefault("FLASH_ATTENTION_TRITON_AMD_ENABLE", "TRUE")
try:
try:
from flash_attn import flash_attn_func # type: ignore
except ImportError:
from flash_attn.flash_attn_interface import flash_attn_func # type: ignore
FLASH_ATTN_AVAILABLE = True
except ImportError:
flash_attn_func = None
FLASH_ATTN_AVAILABLE = False
def get_attn_backend() -> str:
"""Return which attention backend will be used."""
if FLASH_ATTN_AVAILABLE:
return "flash-attn v2 (Dao-AILab / ROCm Triton)"
if torch.cuda.is_available():
device_label = "ROCm AOTriton" if _is_rocm() else "CUDA FlashAttention"
return f"PyTorch SDPA (Native {device_label} Kernel)"
return "PyTorch SDPA (CPU Kernel)"
# -- Unified attention function ---------------------------------------------
def bear_attention(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
n_rep: int = 1,
dropout_p: float = 0.0,
causal: bool = True,
mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Compute scaled dot-product attention with automatic backend selection.
Args:
q: (B, T, n_heads, head_dim) — query, already RoPE'd
k: (B, T, n_kv_heads, head_dim) — key, already RoPE'd
v: (B, T, n_kv_heads, head_dim) — value
n_rep: repeat factor for GQA (n_heads // n_kv_heads)
dropout_p: attention dropout probability
causal: use causal (autoregressive) masking
mask: explicit attention mask (only used in SDPA path, ignored by flash-attn)
Returns:
(B, T, n_heads, head_dim) attention output
"""
if FLASH_ATTN_AVAILABLE and q.is_cuda:
# flash_attn_func takes (B, T, H, D) and handles GQA natively
return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal)
# SDPA path: transpose (B, T, H, D) -> (B, H, T, D)
q = q.transpose(1, 2) # (B, H, T, D)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# PyTorch 2.5+ SDPA supports enable_gqa parameter directly
enable_gqa = (n_rep > 1)
try:
# Try native GQA parameter in SDPA (PyTorch 2.5+)
out = F.scaled_dot_product_attention(
q.contiguous(), k.contiguous(), v.contiguous(),
attn_mask=mask,
is_causal=(causal and mask is None),
dropout_p=dropout_p,
enable_gqa=enable_gqa,
)
except TypeError:
# Fallback for PyTorch versions where enable_gqa isn't available
if n_rep > 1:
k = k.repeat_interleave(n_rep, dim=1)
v = v.repeat_interleave(n_rep, dim=1)
out = F.scaled_dot_product_attention(
q, k, v,
attn_mask=mask,
is_causal=(causal and mask is None),
dropout_p=dropout_p,
)
return out.transpose(1, 2) # back to (B, T, H, D)
|