Spaces:
Running on Zero
Running on Zero
File size: 2,890 Bytes
5c86cdc 8fc8d44 1b5453a 8fc8d44 1b5453a 8fc8d44 5c86cdc 8fc8d44 5c86cdc | 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 | import torch
from einops import rearrange
from torch import Tensor
from typing import Optional, List
import os
# ============================================================
# Flash Attention 3 Support (for Hopper GPUs: H100/H200)
# ============================================================
# Default ON for ZeroGPU (H200), set USE_FA3=0 to disable
_USE_FA3 = os.environ.get("USE_FA3", "1") == "1"
_flash_attn_func = None
if _USE_FA3:
try:
from kernels import get_kernel
_fa3_kernel = get_kernel("kernels-community/vllm-flash-attn3")
_flash_attn_func_raw = _fa3_kernel.flash_attn_func
@torch.library.custom_op("flash::flash_attn_func", mutates_args=())
def _flash_attn_func(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
softmax_scale: Optional[float] = None,
causal: bool = False,
) -> torch.Tensor:
outputs = _flash_attn_func_raw(q, k, v, softmax_scale=softmax_scale, causal=causal)
return outputs[0]
@_flash_attn_func.register_fake
def _(q, k, v, **kwargs):
return torch.empty_like(q).contiguous()
print("✓ Flash Attention 3 loaded successfully!")
except Exception as e:
print(f"Flash Attention 3 not available: {e}")
_USE_FA3 = False
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
q, k = apply_rope(q, k, pe)
if _USE_FA3 and _flash_attn_func is not None:
# FA3 expects (B, L, H, D) format
q_fa3 = rearrange(q, "B H L D -> B L H D")
k_fa3 = rearrange(k, "B H L D -> B L H D")
v_fa3 = rearrange(v, "B H L D -> B L H D")
x = _flash_attn_func(q_fa3, k_fa3, v_fa3)
x = rearrange(x, "B L H D -> B L (H D)")
else:
# Standard PyTorch SDPA (uses FA2 if available)
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
x = rearrange(x, "B H L D -> B L (H D)")
return x
def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
assert dim % 2 == 0
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
omega = 1.0 / (theta**scale)
out = torch.einsum("...n,d->...nd", pos, omega)
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
return out.float()
def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
|