Attention variants

Deep dive

How it works

In: Q [T,H,d_h], K and V [T,H_kv,d_h], a mask · Out: attended values [T,H,d_h] · State: KV cache of 2·layers·H_kv·d_h elements per token · Cost: prefill O(T²·d), decode O(T·d) per token

Shapes for one sequence of length T: Q: [T, H, d_h], K, V: [T, H_kv, d_h], with H % H_kv == 0.

1. Query head h reads KV head h / (H / H_kv). MHA is H_kv = H, MQA is H_kv = 1, GQA is in between.
2. Score s_ij = q_i · k_j / √d_h for allowed j. Causal: j ≤ i. Sliding window w: i − w < j ≤ i.
3. p_ij = exp(s_ij − max_j s_ij) / Σ_j exp(…), output o_i = Σ_j p_ij v_j.
4. Concatenate heads and project with W_o.

MHA   q0 q1 q2 q3 q4 q5 q6 q7      GQA (H_kv=2)  q0..q3 → kv0   q4..q7 → kv1
      k0 k1 k2 k3 k4 k5 k6 k7      MQA (H_kv=1)  q0..q7 → kv0
sliding window w=3, causal mask for query 5:  keys 3 4 5 visible, 0-2 masked

MLA (DeepSeek-V2). Cache one latent c_t = W_dkv·x_t (d_c = 512) per token and layer, plus a shared 64-dim RoPE key. K and V for all heads are up-projections of c_t. The up-projection for K can be absorbed into W_q, and the one for V into W_o, so decode attends directly over the latents. Cache: 576 elements per token per layer, vs 2 · 128 heads · 128 = 32768 for MHA on the same model.

Code

/// Causal grouped-query attention for one sequence, row-major:
/// q: [t][h][d], k, v: [t][h_kv][d]. window = None means full causal attention.
pub fn gqa_attention(q: &[f32], k: &[f32], v: &[f32], t: usize, h: usize, h_kv: usize, d: usize,
                     window: Option<usize>) -> Vec<f32> {
    assert!(h % h_kv == 0);
    let group = h / h_kv;
    let scale = 1.0 / (d as f32).sqrt();
    let mut out = vec![0.0f32; t * h * d];
    let mut s = vec![0.0f32; t];
    for i in 0..t {
        let lo = window.map_or(0, |w| (i + 1).saturating_sub(w)); // first visible key
        for hh in 0..h {
            let kvh = hh / group;                                    // shared KV head
            let qi = &q[(i * h + hh) * d..][..d];
            let mut m = f32::NEG_INFINITY;
            for j in lo..=i {
                let kj = &k[(j * h_kv + kvh) * d..][..d];
                s[j] = qi.iter().zip(kj).map(|(a, b)| a * b).sum::<f32>() * scale;
                m = m.max(s[j]);
            }
            let mut denom = 0.0f32;
            for j in lo..=i { s[j] = (s[j] - m).exp(); denom += s[j]; }
            let o = &mut out[(i * h + hh) * d..][..d];
            for j in lo..=i {
                let p = s[j] / denom;
                let vj = &v[(j * h_kv + kvh) * d..][..d];
                for c in 0..d { o[c] += p * vj[c]; }
            }
        }
    }
    out
}

/// Bytes of KV cache per token (K and V, all layers).
pub fn kv_bytes_per_token(layers: usize, h_kv: usize, d: usize, bytes: usize) -> usize {
    2 * layers * h_kv * d * bytes
}

/// MHA → GQA conversion: mean-pool each group of K (or V) projection heads. w: [h][d][d_model].
pub fn mean_pool_heads(w: &[f32], h: usize, h_kv: usize, per_head: usize) -> Vec<f32> {
    let g = h / h_kv;
    let mut out = vec![0.0f32; h_kv * per_head];
    for hh in 0..h {
        for i in 0..per_head { out[(hh / g) * per_head + i] += w[hh * per_head + i] / g as f32; }
    }
    out
}

fn main() {
    let (t, h, d) = (6, 4, 8);
    let f = |n: usize, s: usize| -> Vec<f32> { (0..n).map(|i| (((i * s) % 17) as f32 - 8.0) / 8.0).collect() };
    let q = f(t * h * d, 7);
    let (k2, v2) = (f(t * 2 * d, 5), f(t * 2 * d, 3));
    // GQA with 2 KV heads == MHA where each KV head is duplicated for its 2 query heads.
    let dup = |x: &[f32]| -> Vec<f32> {
        (0..t * h * d).map(|i| { let (tok, hh, c) = (i / (h * d), (i / d) % h, i % d);
                                 x[(tok * 2 + hh / 2) * d + c] }).collect() };
    let a = gqa_attention(&q, &k2, &v2, t, h, 2, d, None);
    let b = gqa_attention(&q, &dup(&k2), &dup(&v2), t, h, h, d, None);
    assert!(a.iter().zip(&b).all(|(x, y)| (x - y).abs() < 1e-6));
    // A window at least as long as the sequence is plain causal attention.
    let c = gqa_attention(&q, &k2, &v2, t, h, 2, d, Some(t));
    assert!(a.iter().zip(&c).all(|(x, y)| (x - y).abs() < 1e-6));
    let mha = kv_bytes_per_token(80, 64, 128, 2);
    let gqa = kv_bytes_per_token(80, 8, 128, 2);
    println!("Llama-2-70B KV/token: MHA {} KiB, GQA {} KiB; 4k ctx GQA {:.2} GiB",
             mha / 1024, gqa / 1024, (gqa * 4096) as f64 / (1u64 << 30) as f64);
    let pooled = mean_pool_heads(&[1.0, 3.0, 5.0, 7.0], 4, 2, 1);
    assert_eq!(pooled, vec![2.0, 6.0]);
}

Complexity and constants

• Prefill: 4·T²·d_h·H FLOPs per layer for QKᵀ and PV (half with a causal-aware kernel). Window w: 4·T·w·d_h·H.
• Decode, per token and layer: read 2·T·H_kv·d_h cached elements, do 4·T·H·d_h FLOPs. Arithmetic intensity ≈ H/H_kv FLOPs per element read: 1 for MHA, 8 for GQA with group 8. Still far below the hundreds of FLOPs/byte a GPU needs to be compute-bound, so decode attention is bandwidth-bound.
• Llama 3 8B (32 layers, 8 KV heads, d_h 128): 128 KiB/token in bf16, so the 64 GB left on an 80 GB GPU after 16 GB of weights holds about 490k cached tokens across the batch.
• Sliding window w caps the cache at w tokens per layer (a rolling buffer, index pos % w).

Efficiency

• Cost: prefill is compute-bound at 4·T²·d_h·H FLOPs per layer. Decode reads the whole KV cache per token (2·T·H_kv·d_h elements), so it is bandwidth-bound. Llama 3 8B holds 128 KiB of cache per token (this page): at T = 8192 that is 1 GiB per sequence, about 0.3 ms per token to stream at the 3.35 TB/s HBM3 bandwidth of an H100 SXM (arithmetic from the datasheet figure).
• Levers:
  • GQA/MQA — cache and decode traffic shrink by H/H_kv, 8x at group size 8 (Ainslie 2023) — needs training or uptraining.
  • MLA — 93.3% smaller KV cache, 5.76x peak generation throughput vs DeepSeek 67B (DeepSeek-V2) — model trained with it.
  • FlashAttention — no T×T matrix in HBM; up to 3x on GPT-2 at length 1K (Dao 2022) — prefill and training.
  • Paged KV cache — existing systems wasted 60-80% of cache memory, PagedAttention under 4%, 2-4x throughput (Kwon 2023) — many concurrent requests.
• Naive → standard → state of the art:
  • MHA, materialized T×T scores, contiguous cache · O(T²) HBM traffic, cache H times larger
  • GQA (8 KV heads) with a fused flash kernel · cache /8, no T² traffic
  • MLA or GQA plus paged cache and window layers · 93% smaller cache (DeepSeek-V2)

Implementation notes

• Never materialize the T×T score matrix in production. Use a fused kernel (FlashAttention) that streams K/V tiles with an online softmax. The code above is the reference semantics.
• GQA in kernels: load each K/V tile once and reuse it for all H/H_kv query heads in the group. That reuse is the speedup.
• Store the cache in pages (paged KV cache) so memory grows per token instead of per max length.
• Apply RoPE to q and k before caching k. MLA needs a separate small RoPE key, because rotation does not commute with the absorbed up-projection.
• Logit growth: QK-norm (normalization) or soft-capping (c·tanh(s/c), Gemma 2 uses c = 50 for attention logits) keeps scores bounded in bf16.
• Sliding-window serving with sink tokens: keep the first 4 tokens plus the window (StreamingLLM). The softmax dumps excess mass on early tokens and breaks if they are evicted.

Pitfalls

• Wrong head mapping: h % h_kv instead of h / group silently pairs query heads with the wrong KV head. The model runs and is just worse.
• Masking with -inf in fp16 then subtracting the row max on a fully masked row: NaN. Use a large finite negative or skip empty rows.
• Scaling by 1/√d_model instead of 1/√d_h.
• Sliding-window cache indexed by absolute position without the modulo, or RoPE recomputed from the buffer slot instead of the real position.
• Assuming GQA saves FLOPs: it saves memory and bandwidth. Prefill FLOPs are unchanged.

Variants and alternatives

In production

Llama 2 70B and Llama 3 (GQA, 8 KV heads), Mistral 7B (GQA plus sliding window), PaLM (MQA), Gemma 2 (local/global interleave), DeepSeek-V2 and V3 (MLA); Swin Transformer (window attention for image classification and detection). Kernels: FlashAttention 2/3, FlashInfer, vLLM PagedAttention, torch.nn.functional.scaled_dot_product_attention (has enable_gqa in recent versions).

Implementations

Sources

• Liu et al., Swin Transformer: Hierarchical Vision Transformer using Shifted Windows, ICCV 2021. paper · code
• Vaswani et al., Attention Is All You Need, NeurIPS 2017. paper
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019. paper
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023. paper
• Jiang et al., Mistral 7B, 2023. paper · code
• Beltagy, Peters, Cohan, Longformer: The Long-Document Transformer, 2020. paper
• DeepSeek-AI, DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, 2024. paper
• Xiao et al., Efficient Streaming Language Models with Attention Sinks, ICLR 2024. paper
• Gemma Team, Gemma 2: Improving Open Language Models at a Practical Size, 2024. paper
• Touvron et al., Llama 2: Open Foundation and Fine-Tuned Chat Models, 2023. paper
• Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022. paper · code
• Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023. paper · code