Flash Attention 3 and Multi-Headed Latent Attention: The Evolution of Efficient Attention

Attention computation grows O(N²) with context length:

Context Length (N)  Attention FLOPs (d=128)  HBM Access (KV cache)
       4K                   4.2B                     8 MB (MHA)
       8K                   16.8B                    32 MB
       32K                  268B                     512 MB
       128K                 4.3T                     8 GB
       1M                   268T                     512 GB ← impossible

Two orthogonal solutions:
1. FA3: Make each FLOP faster via better GPU utilization (2x speedup)
2. MLA: Reduce KV cache size by 75-93% (4-16x memory reduction)
→ Combined: 8-32x effective improvement

Year  │ Technique                 │ Speedup vs MHA  │ Key Innovation
──────┼───────────────────────────┼─────────────────┼────────────────────
2022  │ Flash Attention v1        │ 2-4x            │ Tiling, online softmax
2023  │ Flash Attention v2        │ 2x vs FA1       │ Reduced non-matmul, better parallelism
2023  │ PageAttention (vLLM)      │ Memory saving   │ Non-contiguous KV cache
2024  │ Flash Attention 3         │ 2x vs FA2       │ WGMMA, FP8, async
2024  │ MLA (DeepSeek)            │ 4-16x KV cache  │ Low-rank KV compression
2024  │ GQA/MQA                   │ 4-8x KV cache   │ Shared KV heads
2025  │ FA3 + MLA hybrid          │ 4-32x eff.      │ Combined
2026  │ Hardware attention (NVL)  │ 10x vs FA3      │ In-memory compute

Feature                     FA2 (A100)          FA3 (H100)
──────────────────────────────────────────────────────────────
Tensor Core generation      MMA (FP16)          WGMMA (FP16/FP8)
Warp scheduling             Manual              Async warp groups
SM-to-SM communication      N/A                 Distributed shared memory
Tile quantization           None                Per-tile FP8
HBM/compute overlap         Minimal             Full overlap
Peak utilization            45% of A100 peak    75% of H100 peak
Supported precision         FP16/BF16           FP16/FP8/INT8

Flash Attention 3: Warp-specialized pipeline

┌─────────────────────────────────────────────────────────┐
│  Warp Group 1: Data Loaders                               │
│  ┌──────┐  ┌──────┐  ┌──────┐                             │
│  │ Q,K,V│  │ Q,K,V│  │ Q,K,V│                             │
│  │ Load │  │ Load │  │ Load │                             │
│  └──┬───┘  └──┬───┘  └──┬───┘                             │
│     │         │         │                                  │
│     ▼         ▼         ▼                                  │
│  ┌──────────────────────────────────────┐                  │
│  │  Distributed Shared Memory (32 MB)    │                 │
│  │  (SM-to-SM communication via DS)      │                 │
│  └──────────────────────────────────────┘                  │
│             │                                              │
│             ▼                                              │
│  Warp Group 2: Compute                                      │
│  ┌──────┐  ┌──────┐  ┌──────┐                             │
│  │WGMMA │  │WGMMA │  │WGMMA │                             │
│  │Compute│  │Compute│  │Compute│                             │
│  └──┬───┘  └──┬───┘  └──┬───┘                             │
│     │         │         │                                  │
│     ▼         ▼         ▼                                  │
│  ┌──────────────────────────────────────┐                  │
│  │  Online Softmax + Output Accumulate   │                 │
│  └──────────────────────────────────────┘                  │
│             │                                              │
│             ▼                                              │
│  Warp Group 3: Store                                        │
│  ┌──────────────────────────────┐                          │
│  │  Write output to HBM         │                          │
│  └──────────────────────────────┘                          │
│                                                             │
│  All three warp groups run concurrently —                    │
│  Data loads, computes, and stores OVERLAP in time           │
└─────────────────────────────────────────────────────────────┘

Traditional MMA (Ampere, FA2):
  Each warp (32 threads): issues MMA instruction independently
  → 4 warps = 4 independent MMA instructions
  → Requires manual data distribution across warps
  → Significant setup overhead

WGMMA (Hopper, FA3):
  Warp group (128 threads): issues ONE WGMMA instruction
  → Tensor Cores auto-distribute across all 128 threads
  → Hardware handles data routing
  → 50% less instruction overhead
  → Higher Tensor Core utilization

Performance impact:
  MMA: 312 TFLOPS peak (H100, FP16, theoretical)
  WGMMA: 396 TFLOPS peak (H100, FP16, theoretical)
  → 27% higher Tensor Core throughput

# Conceptual pseudocode for FA3's WGMMA usage
# (actual implementation is in CUDA PTX inline assembly)

@triton.jit
def fa3_attention_kernel_hopper(
    q_ptr, k_ptr, v_ptr, o_ptr,
    N, d, num_heads,
    stride_q, stride_k, stride_v, stride_o,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_DMODEL: tl.constexpr,
):
    """Flash Attention 3 using WGMMA and async pipeline."""

    pid = tl.program_id(0)

    # Use Hopper's warp-specialized approach:
    # Split warps into: (1) data movers, (2) compute

    # Phase 1: Load Q tile via async copy (data mover warps)
    q_tile = tl.load(
        q_ptr + pid * BLOCK_SIZE_M * stride_q + 
        tl.arange(0, BLOCK_SIZE_M)[:, None] * stride_q +
        tl.arange(0, BLOCK_DMODEL)[None, :]
    )

    # Phase 2: WGMMA-based block attention
    # WGMMA loads from shared memory (already prefetched)
    acc = tl.zeros([BLOCK_SIZE_M, BLOCK_DMODEL], dtype=tl.float32)
    m_i = tl.full([BLOCK_SIZE_M], -float('inf'), dtype=tl.float32)
    l_i = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32)

    for start_n in range(0, N, BLOCK_SIZE_N):
        # Prefetch K,V tiles into shared memory (async)
        # Using Hopper's cp.async (fully overlapped with computation)

        k_tile = tl.load(...)   # Async load
        v_tile = tl.load(...)   # Async load

        # WGMMA: Matrix multiply using Tensor Cores
        # Unlike MMA in FA2, WGMMA handles all 4 warps' work
        # in a single instruction with auto-routing

        s = tl.dot(q_tile, k_tile.T, use_wgmma=True)
        s *= (1.0 / tl.sqrt(tl.cast(BLOCK_DMODEL, tl.float32)))

        # Online softmax (same as FA2)
        # ... (standard online softmax)

        # WGMMA: Accumulate V
        acc = tl.dot(exp_s.to(tl.float16), v_tile, acc=acc, use_wgmma=True)

    # Normalize and store
    acc = acc / l_i[:, None]
    tl.store(o_ptr + ... , acc)

FA2 Timeline (sequential):
Time → [Load Q] [Load KV] [Compute S] [Softmax] [Compute O] [Store O] [Load next KV]
        └───────── HBM idle during compute ─────────┘
        HBM utilization: ~45%

FA3 Timeline (overlapped):
Time → [Load Q] [Compute S] [Softmax] [Compute O] [Store O]
                  │  Async load KV₁   │
                  │  Async load KV₂   │
                  │  Async load KV₃   │
                  └─── Fully overlapped ──┘
        ┌──────────────────────────────┐
        HBM utilization: ~75-80%

class AsyncPipeliner:
    """Asynchronous pipeline for FA3's HBM/compute overlap."""

    def __init__(self, gpu):
        self.copy_engine = gpu.async_copy_engine  # H100's async copy
        self.compute = gpu.wgmma_compute

    def attention_with_overlap(self, Q, K, V):
        """Fully overlapped attention computation."""

        # Number of KV blocks to prefetch
        N_PREFETCH = 2  # Prefetch 2 blocks ahead

        # Initialize async copies
        prefetch_buffer = [
            self.copy_engine.async_read(K[0]),  # Start loading K₀
            self.copy_engine.async_read(V[0]),  # Start loading V₀
        ]

        for i in range(len(K)):
            # Wait for current block's data
            prefetch_buffer[0].wait()
            prefetch_buffer[1].wait()

            # Start loading NEXT block(s)
            next_idx = i + N_PREFETCH
            if next_idx < len(K):
                prefetch_buffer[0] = self.copy_engine.async_read(K[next_idx])
                prefetch_buffer[1] = self.copy_engine.async_read(V[next_idx])

            # Compute attention for current block
            # This runs WHILE next block is loading in background
            S = self.compute.matmul(Q, K[i].T)
            S = self.online_softmax(S)
            acc = self.compute.matmul(S, V[i])

        return acc

Hopper SM has 4 warp schedulers → run 4 warps concurrently

FA3 warp specialization:

Warp 0-1 (Data Movers):
  - Issue cp.async to load tiles from HBM → shared memory
  - Manage prefetch buffer
  - Execute async store for completed output tiles

Warp 2-3 (Compute):
  - Execute WGMMA for matmul
  - Compute online softmax
  - Accumulate attention output
  - NEVER wait on HBM (data is always ready)

SM resources:
  - 228 KB shared memory → ~170 KB for tiles, ~58 KB for buffers
  - 65K registers → ~16K per warp for WGMMA operands
  - 4 warp schedulers → 2 data, 2 compute (balanced)

Key metric: 
  FA2: 45% of H100 peak TFLOPS (idle warps waiting for data)
  FA3: 75% of H100 peak TFLOPS (all warps always busy)

Standard FP16 attention:
  q_fp16 × k_fp16^T = s_fp16 → softmax(s_fp16) = p_fp16 → p_fp16 × v_fp16

FP8 attention (FA3):
  q_fp16 → quantize to FP8 (per-tile scale) → q_fp8
  k_fp16 → quantize to FP8 (per-tile scale) → k_fp8
  q_fp8 × k_fp8^T = s_fp16 → (dequantized and accumulated in FP16)

  softmax(s_fp16) = p_fp16 → quantize to FP8 → p_fp8
  v_fp16 → quantize to FP8 → v_fp8
  p_fp8 × v_fp8 = o_fp16 → (dequantized and accumulated in FP16)

Key insight: Attention inner loop is bandwidth-bound, not compute-bound
→ FP8 Tensor Cores give 2x throughput vs FP16
→ Only ~0.1% accuracy loss at tile level
→ Large accuracy gains from longer context (FA3 enables 2x larger tiles)

class FA3FP8BlockProcessor:
    """FP8 block processing with per-tile quantization."""

    def process_block(self, Q_fp16, K_fp16, V_fp16):
        """Process one attention block in FP8."""

        # 1. Per-tile quantization to FP8
        # Compute scales from absolute max (E4M3 format)
        q_scale = Q_fp16.abs().max() / 448.0  # FP8 E4M3 max = 448
        k_scale = K_fp16.abs().max() / 448.0
        v_scale = V_fp16.abs().max() / 448.0

        Q_fp8 = (Q_fp16 / q_scale).to(torch.float8_e4m3fn)
        K_fp8 = (K_fp16 / k_scale).to(torch.float8_e4m3fn)
        V_fp8 = (V_fp16 / v_scale).to(torch.float8_e4m3fn)

        # 2. Matrix multiply in FP8 using Tensor Cores (2x throughput)
        S_fp8 = torch._scaled_mm(
            Q_fp8, K_fp8.t(),
            scale_a=1.0, scale_b=1.0,
            out_dtype=torch.float16
        )

        # 3. Dequantize attention scores
        S_fp16 = S_fp8.float() * (q_scale * k_scale)
        S_fp16 /= math.sqrt(Q_fp16.shape[-1])  # Scale

        # 4. Softmax in FP16 (small portion of total compute)
        P_fp16 = torch.softmax(S_fp16, dim=-1)

        # 5. Quantize P for FP8 matmul with V
        p_scale = P_fp16.abs().max() / 448.0
        P_fp8 = (P_fp16 / p_scale).to(torch.float8_e4m3fn)

        # 6. FP8 matmul with V (2x throughput)
        O_fp8 = torch._scaled_mm(P_fp8, V_fp8, scale_a=p_scale, scale_b=v_scale)

        return O_fp8

Standard MHA KV cache per token (DeepSeek-V2, 64 heads, 128d):
  Keys:   64 heads × 128d × 2 bytes = 16,384 bytes/token
  Values: 64 heads × 128d × 2 bytes = 16,384 bytes/token  
  Total:  32,768 bytes/token (32 KB/token)

For 128K context: 128K × 32 KB = 4 GB (per sequence!)
For batch=64, 128K context: 256 GB (impossible on single GPU)

GQA (8 groups): 128K × 4 KB = 512 MB/seq → Better, but still large
MLA:            128K × 2 KB = 256 MB/seq → 16x vs MHA!

Standard MHA (for comparison):
  Input: h (hidden state)
    K = h × W_k       → [num_heads × head_dim]
    V = h × W_v       → [num_heads × head_dim]
  KV cache: K, V for each token = 2 × num_heads × head_dim

MLA (DeepSeek-V2/V3):
  Input: h (hidden state)
    c = h × W_dkv     → [d_c]  (compressed latent, d_c << 2 × num_heads × head_dim)

    K = c × W_uk      → [num_heads × head_dim]  (up-projected from latent)
    V = c × W_uv      → [num_heads × head_dim]  (up-projected from latent)

  KV cache: c only (shared latent) → [d_c]
    Plus small per-head bias vectors stored in model weights (not cached!)

  d_c = 512 (DeepSeek-V2) vs 2 × 64 × 128 = 16,384 (MHA)
  Compression ratio: 16384/512 = 32x!

But wait—the up-projection matrices W_uk, W_uv are learned weights.
K and V are computed ON THE FLY from c during attention.
The KV cache stores only c (compressed latent).

┌────────────────────────────────────────────────────────────┐
│  Multi-Headed Latent Attention Architecture                 │
├────────────────────────────────────────────────────────────┤
│                                                              │
│  Input: h (hidden_dim = 4096)                               │
│    │                                                        │
│    ▼                                                        │
│  h × W_dkv → c (compressed latent, dim = 512)              │
│    │                                                        │
│    │  KV Cache stores ONLY c (512 floats = 1KB/token)      │
│    │                                                        │
│    ├──→ c × W_uk → K (64 heads × 128d = 8192 dims)        │
│    │    (up-projected on-the-fly during attention)          │
│    │                                                        │
│    └──→ c × W_uv → V (64 heads × 128d = 8192 dims)        │
│         (up-projected on-the-fly during attention)          │
│                                                              │
│  Similarly for Query:                                        │
│  h × W_q → q (64 heads × 128d)                              │
│    │                                                        │
│    (No caching needed for q—only 1 token at a time)         │
│                                                              │
│  KV cache memory (128K tokens per sequence):                │
│    MHA:   128K × 16KB = 2GB  (64 heads, FP16)              │
│    MQA:   128K × 0.5KB = 64MB  (1 KV head)                 │
│    GQA:   128K × 2KB = 256MB  (8 KV groups)                │
│    MLA:   128K × 1KB = 128MB  (compressed latent, 512d)    │
│                                                              │
│  MLA advantage grows with context length:                   │
│    4K: MHA=64MB, MLA=4MB (16x)                              │
│    128K: MHA=2GB, MLA=128MB (16x)                           │
│    1M: MHA=16GB, MLA=1GB (16x) ← fits on consumer GPU!     │
│                                                              │
└──────────────────────────────────────────────────────────────┘

Standard MHA:
  K = h · W_k         ∈ ℝ^(n_heads × d_head)
  V = h · W_v         ∈ ℝ^(n_heads × d_head)

  KV_cache[t] = concat(K, V)  ∈ ℝ^(2 × n_heads × d_head)

MLA with Low-Rank Compression:
  c_KV = h · W_dkv    ∈ ℝ^(d_c)     # Compressed latent (shared for K and V)
  K_up = c_KV · W_uk  ∈ ℝ^(n_heads × d_head)   # Up-projected keys
  V_up = c_KV · W_uv  ∈ ℝ^(n_heads × d_head)   # Up-projected values

  KV_cache[t] = c_KV  ∈ ℝ^(d_c)     # Only store the latent!

Where:
  d_c << 2 × n_heads × d_head
  (512 vs 16,384 for 64 heads × 128d)

Additionally, MLA uses:
- Per-head bias vectors b_k, b_v (absorbed into up-projection)
- Rotary position embeddings applied AFTER compression
  → Can rotate keys efficiently on-the-fly

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

class MultiHeadLatentAttention(nn.Module):
    """Multi-Headed Latent Attention from DeepSeek-V2/V3."""

    def __init__(self, d_model, n_heads, d_head, d_c):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_head
        self.d_c = d_c  # Compressed latent dimension

        # Query projection (standard)
        self.w_q = nn.Linear(d_model, n_heads * d_head, bias=False)

        # Compressed KV latent projection (shared)
        self.w_dkv = nn.Linear(d_model, d_c, bias=False)  # Down-project

        # Up-projection from latent to full K, V
        self.w_uk = nn.Linear(d_c, n_heads * d_head, bias=False)
        self.w_uv = nn.Linear(d_c, n_heads * d_head, bias=False)

        # Output projection
        self.w_o = nn.Linear(n_heads * d_head, d_model, bias=False)

        # KV cache: store only compressed latents
        self.kv_cache = None  # [batch, seq_len, d_c]

    def forward(self, x, use_cache=False):
        batch, seq_len, _ = x.shape

        # Q projection
        q = self.w_q(x)  # [batch, seq_len, n_heads * d_head]
        q = q.view(batch, seq_len, self.n_heads, self.d_head)

        # Compressed KV latent (what gets cached!)
        c_kv = self.w_dkv(x)  # [batch, seq_len, d_c]

        if use_cache and self.kv_cache is not None:
            # Append to cache (only compressed latent!)
            self.kv_cache = torch.cat([self.kv_cache, c_kv], dim=1)
            full_seq_len = self.kv_cache.shape[1]

            # Up-project ALL cached latents to full K, V
            c_all = self.kv_cache  # [batch, full_seq_len, d_c]
            k = self.w_uk(c_all).view(batch, full_seq_len, self.n_heads, self.d_head)
            v = self.w_uv(c_all).view(batch, full_seq_len, self.n_heads, self.d_head)
        else:
            # First forward pass — compute from scratch
            k = self.w_uk(c_kv).view(batch, seq_len, self.n_heads, self.d_head)
            v = self.w_uv(c_kv).view(batch, seq_len, self.n_heads, self.d_head)

            if use_cache:
                self.kv_cache = c_kv

        # Standard attention (but K, V are up-projected from latent)
        # Apply RoPE here (omitted for clarity)
        attn_output = F.scaled_dot_product_attention(
            q.transpose(1, 2),  # [batch, n_heads, seq_len, d_head]
            k.transpose(1, 2),
            v.transpose(1, 2),
        )

        # Concatenate heads and project output
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch, seq_len, -1)
        output = self.w_o(attn_output)

        return output

    def clear_cache(self):
        self.kv_cache = None

Prefill phase (process entire prompt at once):
  - Standard: Compute K, V for all tokens → store in cache
  - MLA: Compute c_KV for all tokens → up-project to K, V → attention
  - MLA overhead: extra W_uk, W_uv matmuls (d_c × n_heads × d_head)
  - Cost: ~10-15% more FLOPs during prefill
  - Benefit: 16x smaller KV cache → can handle 16x longer prompts

Decode phase (generate one token at a time):
  - Standard: 
    Load KV cache from HBM: 2 × n_heads × d_head × 2 bytes/token
    Compute single-token attention: O(n_heads × d_head × seq_len)
  - MLA:
    Load c_KV cache from HBM: d_c × 2 bytes/token (16x less)
    Up-project c → K, V: matmul (d_c → n_heads × d_head)
    Compute single-token attention: same as standard
  - Benefit: 16x less HBM bandwidth → 2-3x faster decoding

def estimate_decode_speed(model_type, seq_len, batch_size, gpu_bw=3350):
    """Estimate decode throughput (tok/s/GPU) for MHA vs MLA."""

    models = {
        "deepseek_v2_mha": {"n_heads": 64, "d_head": 128, "d_model": 4096},
        "deepseek_v2_mla": {"n_heads": 64, "d_head": 128, "d_model": 4096, "d_c": 512},
    }

    results = {}

    for name, config in models.items():
        is_mla = "d_c" in config

        if is_mla:
            # MLA: load compressed latent (d_c floats)
            kv_bytes_per_token = config["d_c"] * 2
        else:
            # MHA: load full K, V
            kv_bytes_per_token = 2 * config["n_heads"] * config["d_head"] * 2

        # Total KV cache read per step
        total_kv_bytes = kv_bytes_per_token * seq_len * batch_size

        # Attention compute (same for both)
        compute_flops = 2 * batch_size * config["n_heads"] * config["d_head"] * seq_len

        # MLA has extra up-projection matmul
        if is_mla:
            compute_flops += 2 * batch_size * config["d_c"] * config["n_heads"] * config["d_head"]

        # Memory-bound: most time is HBM read
        time_hbm = total_kv_bytes / (gpu_bw * 1e9)  # seconds
        time_compute = compute_flops / (989 * 1e12)  # seconds (H100 FP16 peak)

        # Actual: dominated by HBM
        time_per_step = max(time_hbm, time_compute)
        tokens_per_sec = batch_size / time_per_step

        results[name] = {
            "kv_bytes_per_step": total_kv_bytes,
            "time_hbm_ms": time_hbm * 1000,
            "time_compute_ms": time_compute * 1000,
            "tokens_per_sec": tokens_per_sec,
        }

    return results

# Example: 128K context, batch=1
# MHA:  ~15 tok/s (HBM bound: 4GB read per step)
# MLA:  ~45 tok/s (3x faster, 256MB read per step)

┌──────────────────────────────────────────────────────────┐
│  Combined FA3 + MLA Attention Kernel                       │
├──────────────────────────────────────────────────────────┤
│                                                            │
│  1. MLA-compressed KV cache (d_c=512) stored in HBM       │
│                                                            │
│  2. During prefill:                                        │
│     - Load c_KV from HBM (1KB/token)                       │
│     - Up-project using WGMMA on Hopper Tensor Cores       │
│     - Apply FA3 FP8 block processing for attention        │
│     - Store only c_KV back to HBM (not full K,V)          │
│                                                            │
│  3. During decode:                                         │
│     - Load entire c_KV cache: 128K × 1KB = 128MB          │
│     - Up-project all cached latents via batched WGMMA      │
│     - Single FA3 attention pass (FP8, overlapped)          │
│     - Store new c_KV (1KB) to cache                        │
│                                                            │
│  Combined benefit:                                         │
│  - FA3: 2x faster attention (WGMMA + FP8 + overlap)       │
│  - MLA: 3x faster decoding (16x less HBM reads)           │
│  - Total: 6x speedup over standard attention               │
│  - Plus: 16x longer context with same memory               │
│                                                            │
└──────────────────────────────────────────────────────────┘

@triton.jit
def fa3_mla_attention_kernel(
    # MLA compressed cache
    c_kv_ptr,        # [batch, seq_len, d_c]
    w_uk_ptr,        # [d_c, n_heads * d_head] 
    w_uv_ptr,        # [d_c, n_heads * d_head]
    # Query
    q_ptr,           # [batch, n_heads * d_head]
    # Output
    o_ptr,           # [batch, n_heads * d_head]
    # Dimensions
    seq_len, d_c, n_heads, d_head,
    stride_c, stride_q, stride_o,
    BLOCK_DC: tl.constexpr,
    BLOCK_HEADDIM: tl.constexpr,
    BLOCK_SEQ: tl.constexpr,
):
    """Combined FA3 + MLA attention kernel."""

    pid = tl.program_id(0)
    head_id = pid

    # 1. Load MLA compressed KV cache blocks
    # FA3-style async loading with prefetch

    # 2. Up-project c_KV to K, V using WGMMA
    # Essential MLA step: reconstruct full K, V from latent

    # 3. FA3-style attention with the reconstructed K, V
    # Using FP8 WGMMA, online softmax, overlapped

    # 4. Store output

    # (Full implementation combines both techniques)
    pass