for i in range(0, N, B):        # Q blocks
    for j in range(0, N, B):    # K,V blocks
        S[i,j] = Q[i] @ K[j].T
        P[i,j] = softmax(S[i,j])
        O[i] += P[i,j] @ V[j]

# Standard
P = softmax(S_i)

# Online (Flash)
m_i = max(m_i_prev, S_i)
l_i = l_i_prev * exp(m_i_prev - m_i) + sum(exp(S_i - m_i))
P_i = exp(S_i - m_i) / l_i

Input: Q, K, V ∈ R^{N×d}, block_size B
Output: O ∈ R^{N×d}

Initialize:
  O = 0, l = 0, m = -∞   (per row)

for i in 0..N/B:
  Qi = Q[i*B:(i+1)*B]
  Oi = 0, li = 0, mi = -∞

  for j in 0..N/B:
    Kj = K[j*B:(j+1)*B]
    Vj = V[j*B:(j+1)*B]

    Sij = Qi @ Kj.T / sqrt(d)     # (B, B)

    mj = max over rows in Sij
    Pij = exp(Sij - mj)           # (B, B)
    lij = sum over rows in Pij

    # Rescale previous Oi
    scale = exp(mi - mj)
    Oi = Oi * scale
    li = li * scale + lij

    # Update
    Oi = Oi + Pij @ Vj
    mi = max(mi, mj)

  O[i*B:(i+1)*B] = Oi / li

import torch
import torch.nn.functional as F
import math

def flash_attention(Q, K, V, causal=True, block_size=64):
    """
    Q, K, V: (B, N, H, d) or (B, N, d)
    Returns: (B, N, H, d) or (B, N, d)
    """
    if Q.dim() == 3:
        Q, K, V = Q.unsqueeze(2), K.unsqueeze(2), V.unsqueeze(2)

    B, N, H, d = Q.shape
    scale = 1.0 / math.sqrt(d)

    O = torch.zeros_like(Q)
    l = torch.zeros(B, N, H, device=Q.device)
    m = torch.full((B, N, H), -float('inf'), device=Q.device)

    # Precompute causal mask
    if causal:
        mask = torch.triu(torch.ones(N, N, device=Q.device), diagonal=1).bool()

    for i in range(0, N, block_size):
        Qi = Q[:, i:i+block_size]
        end_i = min(i + block_size, N)
        Oi = torch.zeros_like(Qi)
        li = torch.zeros(B, end_i-i, H, device=Q.device)
        mi = torch.full((B, end_i-i, H), -float('inf'), device=Q.device)

        for j in range(0, N, block_size):
            Kj = K[:, j:j+block_size]
            Vj = V[:, j:j+block_size]
            end_j = min(j + block_size, N)

            # Sij: (B, Bi, Bj)
            Sij = torch.matmul(Qi, Kj.transpose(-1, -2)) * scale

            # Causal mask
            if causal:
                mask_ij = mask[i:end_i, j:end_j]
                Sij = Sij.masked_fill(mask_ij, -float('inf'))

            # Online softmax stats
            mj = Sij.max(dim=-1, keepdim=True).values
            Pij = torch.exp(Sij - mj)
            lij = Pij.sum(dim=-1, keepdim=True)

            # Rescale previous
            scale_prev = torch.exp(mi - mj)
            Oi = Oi * scale_prev
            li = li * scale_prev + lij

            # Update output
            Oi = Oi + torch.matmul(Pij, Vj)
            mi = torch.max(mi, mj)

        # Write to global memory
        O[:, i:end_i] = Oi / li
        l[:, i:end_i] = li.squeeze(-1)
        m[:, i:end_i] = mi.squeeze(-1)

    return O.squeeze(2) if O.size(2) == 1 else O

# Generate data
B, N, H, d = 1, 2048, 8, 64
Q = torch.randn(B, N, H, d, device='cuda', dtype=torch.float16)
K = torch.randn_like(Q)
V = torch.randn_like(Q)

# Standard
def standard_attention(Q, K, V):
    scale = 1.0 / math.sqrt(d)
    S = torch.matmul(Q, K.transpose(-1, -2)) * scale
    S = S.masked_fill(torch.triu(torch.ones(N, N, device='cuda'), diagonal=1).bool(), -float('inf'))
    P = F.softmax(S, dim=-1)
    return torch.matmul(P, V)

# Warmup
standard_attention(Q, K, V)
flash_attention(Q, K, V)

# Benchmark
import time
torch.cuda.synchronize()

start = time.time()
for _ in range(10):
    standard_attention(Q, K, V)
torch.cuda.synchronize()
std_time = time.time() - start

start = time.time()
for _ in range(10):
    flash_attention(Q, K, V)
torch.cuda.synchronize()
flash_time = time.time() - start

print(f"Standard: {std_time:.3f}s")
print(f"Flash:    {flash_time:.3f}s")
print(f"Speedup:  {std_time/flash_time:.1f}x")

# Memory
std_mem = torch.cuda.max_memory_allocated() / 1e6
torch.cuda.reset_peak_memory_stats()
flash_attention(Q, K, V)
flash_mem = torch.cuda.max_memory_allocated() / 1e6

print(f"Standard memory: {std_mem:.1f} MB")
print(f"Flash memory:    {flash_mem:.1f} MB")
print(f"Reduction:       {(std_mem - flash_mem)/std_mem*100:.1f}%")

class FlashAttentionBlock(nn.Module):
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.n_head = n_head
        self.Wqkv = nn.Linear(n_embd, 3 * n_embd)
        self.Wo = nn.Linear(n_embd, n_embd)

    def forward(self, x):
        B, T, C = x.shape
        qkv = self.Wqkv(x).reshape(B, T, 3, self.n_head, C // self.n_head)
        q, k, v = qkv.unbind(2)  # (B, T, H, d)
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        out = flash_attention(q, k, v, causal=True)  # (B, H, T, d)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.Wo(out)

import torch
import math

def flash_attention(Q, K, V, causal=True, block_size=64):
    B, N, H, d = Q.shape
    scale = 1.0 / math.sqrt(d)
    O = torch.zeros_like(Q)
    l = torch.zeros(B, N, H, device=Q.device)
    m = torch.full((B, N, H), -float('inf'), device=Q.device)

    if causal:
        mask = torch.triu(torch.ones(N, N, device=Q.device), diagonal=1).bool()

    for i in range(0, N, block_size):
        Qi = Q[:, i:i+block_size]
        end_i = min(i + block_size, N)
        Oi = torch.zeros(B, end_i-i, H, d, device=Q.device)
        li = torch.zeros(B, end_i-i, H, device=Q.device)
        mi = torch.full((B, end_i-i, H), -float('inf'), device=Q.device)

        for j in range(0, N, block_size):
            Kj = K[:, j:j+block_size]
            Vj = V[:, j:j+block_size]
            end_j = min(j + block_size, N)

            Sij = torch.matmul(Qi, Kj.transpose(-1, -2)) * scale
            if causal:
                Sij = Sij.masked_fill(mask[i:end_i, j:end_j], -float('inf'))

            mj = Sij.max(-1, keepdim=True).values
            Pij = torch.exp(Sij - mj)
            lij = Pij.sum(-1, keepdim=True)

            scale_prev = torch.exp(mi - mj)
            Oi = Oi * scale_prev
            li = li * scale_prev + lij
            Oi = Oi + torch.matmul(Pij, Vj)
            mi = torch.max(mi, mj)

        O[:, i:end_i] = Oi / li.unsqueeze(-1)

    return O