S115 · Efficient Attention — Flash, Sparse, Linear

"""attention_bench.py — measure vanilla vs Flash-Attention.
 
Requires torch >= 2.0 (SDPA), a CUDA GPU with compute capability >= 8.0
(Ampere/Hopper) for the Flash backend. Falls back to memory-efficient
on older GPUs.
"""
from __future__ import annotations
import time
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
 
assert torch.cuda.is_available(), "This bench needs a CUDA GPU."
DEVICE, DTYPE = "cuda", torch.float16
 
# ---------- 1. Synthetic Q,K,V tensors ----------
B, H, N, D = 2, 32, 4096, 128       # batch, heads, seq_len, head_dim
Q = torch.randn(B, H, N, D, device=DEVICE, dtype=DTYPE)
K = torch.randn(B, H, N, D, device=DEVICE, dtype=DTYPE)
V = torch.randn(B, H, N, D, device=DEVICE, dtype=DTYPE)
 
 
def bench(fn, warm=3, runs=10) -> float:
    for _ in range(warm):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(runs):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / runs * 1000  # ms per iter
 
 
# ---------- 2. Vanilla (math) backend ----------
def vanilla():
    with sdpa_kernel(SDPBackend.MATH):
        return F.scaled_dot_product_attention(Q, K, V, is_causal=True)
 
# ---------- 3. Flash backend ----------
def flash():
    with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
        return F.scaled_dot_product_attention(Q, K, V, is_causal=True)
 
# ---------- 4. Memory-efficient backend (falls back on older GPUs) ----------
def memeff():
    with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
        return F.scaled_dot_product_attention(Q, K, V, is_causal=True)
 
 
print(f"Sequence: B={B}, H={H}, N={N}, D={D}, dtype={DTYPE}")
try:
    t_math  = bench(vanilla)
    print(f"vanilla MATH        : {t_math:6.2f} ms/iter (baseline)")
except Exception as e:
    print(f"vanilla MATH        : OOM or error — {e.__class__.__name__}")
 
try:
    t_flash = bench(flash)
    print(f"FLASH_ATTENTION     : {t_flash:6.2f} ms/iter"
          f" ({t_math/t_flash:.1f}× faster)")
except Exception as e:
    print(f"FLASH_ATTENTION     : unavailable — {e.__class__.__name__}: {e}")
 
try:
    t_eff = bench(memeff)
    print(f"EFFICIENT_ATTENTION : {t_eff:6.2f} ms/iter"
          f" ({t_math/t_eff:.1f}× faster)")
except Exception as e:
    print(f"EFFICIENT_ATTENTION : unavailable — {e.__class__.__name__}")
 
 
# ---------- 5. Compare peak memory for one call each ----------
def peak_mem(fn) -> float:
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    fn()
    return torch.cuda.max_memory_allocated() / 1024**3  # GB
 
 
print("\n--- Peak memory per attention call ---")
print(f"vanilla MATH        : {peak_mem(vanilla):.2f} GB")
print(f"FLASH_ATTENTION     : {peak_mem(flash):.2f} GB")
print(f"EFFICIENT_ATTENTION : {peak_mem(memeff):.2f} GB")
 
# ---------- 6. HuggingFace model with different attn implementations ----------
from transformers import AutoModelForCausalLM, AutoTokenizer
 
MODEL = "meta-llama/Llama-3.2-1B"    # small enough for a single GPU
try:
    tok = AutoTokenizer.from_pretrained(MODEL)
    for impl in ("eager", "sdpa", "flash_attention_2"):
        try:
            m = AutoModelForCausalLM.from_pretrained(
                MODEL, torch_dtype=DTYPE, attn_implementation=impl
            ).to(DEVICE).eval()
            ids = tok("Hello world " * 500, return_tensors="pt").input_ids.to(DEVICE)
            def gen():
                with torch.no_grad():
                    m.generate(ids, max_new_tokens=32, do_sample=False)
            print(f"\n{impl:<20}: {bench(gen, warm=1, runs=3):.1f} ms/iter")
            del m; torch.cuda.empty_cache()
        except Exception as e:
            print(f"{impl:<20}: unavailable — {e.__class__.__name__}")
except Exception as e:
    print(f"\n(HuggingFace section skipped: {e.__class__.__name__})")