Attention Mechanisms for Generation

Module: Generative AI | Difficulty: Advanced

Standard Attention Complexity

Sparse Attention Patterns

| Pattern | Complexity | Receptive Field | |---------|-----------|-----------------| | Global | | Full | | Local | | Window | | Strided | | Dilated | | Random | | Probabilistic |

Linear Attention (Performer)

where

Flash Attention

Tiling-based exact attention with memory:

def flash_attention_forward(Q, K, V, block_size=256):
    n = Q.size(0)
    O = torch.zeros_like(Q)
    L = torch.zeros(n, 1, device=Q.device)
    M = torch.full((n, 1), float('-inf'), device=Q.device)
    for j in range(0, n, block_size):
        Kj, Vj = K[j:j+block_size], V[j:j+block_size]
        S = Q @ Kj.T / Q.size(-1)**0.5
        M_new = torch.max(M, S.max(dim=-1, keepdim=True).values)
        P = torch.exp(S - M_new)
        L_new = torch.exp(M - M_new) * L + P.sum(dim=-1, keepdim=True)
        O = torch.exp(M - M_new) * O + P @ Vj
        M, L = M_new, L_new
    return O / L

Research Insight: Flash attention is not an approximation — it computes the exact same result as standard attention but with better memory access patterns. The speedup comes from reducing HBM accesses.