flash attn2 EU AI Act Compliance Profile

Flash Attention

Flash Attention is a fast and memory-efficient implementation of the attention mechanism, designed to work with large models and long sequences. This is a Hugging Face compliant kernel build of Flash Attention.

Original code here https://github.com/Dao-AILab/flash-attention.

scripts/readme_example.py provides a simple example of how to use the Flash Attention kernel in PyTorch. It demonstrates standard attention, causal attention, and variable-length sequences.

# /// script
# dependencies = [
#   "numpy", 
#   "torch", 
#   "kernels"
# ]
# ///
import torch
from kernels import get_kernel

# Setup
torch.manual_seed(42)
flash_attn = get_kernel("kernels-community/flash-attn2")
device = torch.device("cuda")

# Create test tensors
B, S, H, D = 2, 5, 4, 8  # batch, seq_len, heads, head_dim
q = k = v = torch.randn(B, S, H, D, device=device, dtype=torch.float16)

# Reference implementation using PyTorch SDPA
def reference_attention(query, key, value, causal=False):
    query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value))
    with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
        out = torch.nn.functional.scaled_dot_product_attention(query, key, value, is_causal=causal)
    return out.transpose(1, 2).contiguous()

# 1. Standard attention
print("\n1. Standard attention:")
out_ref = reference_attention(q, k, v)
out_flash = flash_attn.fwd(
    q=q, 
    k=k, 
    v=v, 
    is_causal=False,
)[0]
print(f"Reference output: {out_ref.shape}")
print(f"Flash output: {out_flash.shape}")
print(f"Outputs close: {torch.allclose(out_flash, out_ref, atol=1e-2, rtol=1e-3)}")

# 2. Causal attention (for autoregressive models)
print("\n2. Causal attention:")

out_ref_causal = reference_attention(q, k, v, causal=True)
out_causal = flash_attn

Your obligations as a flash attn2 deployer

Track your progress

Sign up to track obligations, receive deadline reminders, and generate audit packages.

Tags

Library