Skip to content

Attention variants: MQA, GQA, MLA and FlashAttention

  • Transformers
  • Asked often
  • About 2 h to learn
  • 4 vetted resources
Build a plan that includes this

What this topic covers

Standard multi-head attention keeps a key and value head for every query head, so the KV cache grows fast with context length and batch size. Multi-query and grouped-query attention let several query heads share one K/V head, which shrinks the cache and the bytes read per decoded token. Multi-head latent attention instead caches one small compressed vector per token and expands it into K and V when needed. FlashAttention fixes a different problem: it computes exact attention in tiles held in on-chip SRAM, so the n×n score matrix is never written to HBM. It gets faster by moving less data, not by doing fewer FLOPs. A strong answer says what each technique saves, gives the KV-cache formula with the number of KV heads in it, and can sketch a GQA layer that repeats K/V heads across query groups.

Builds on

Leads to

You’re ready when you can

What interviewers check, from the resources below.

  • Can write the KV-cache size formula with n_kv_heads in it
  • Can place MHA, GQA and MQA on one spectrum of KV groups
  • Can implement GQA by repeating K/V heads across query groups
  • Can explain why whole-model memory savings are smaller than the attention-only estimate
  • Can explain compute-bound vs. memory-bound and why attention is memory-bound
  • Can explain why the n×n score matrix never reaches HBM
  • Can describe online softmax rescaling in two sentences
  • Can contrast what GQA/MLA save (cache size, bytes per decoded token) with what FlashAttention saves (HBM traffic for scores)
  • Can walk through Algorithm 1’s outer and inner loops
  • Can state the HBM access cost of standard attention vs. FlashAttention

What to study, in order

Each link opened and checked by hand; the date says when.

  1. Start here

    Grouped-Query Attention (GQA) (opens in a new tab)

    LLMs-from-scratch (GitHub) · Sebastian Raschka · 45 min · Free · checked 21 Oct 2026

    Why
    Once you mention the KV cache, expect “how does GQA shrink it?” or “write a GQA layer.” This folder puts working MHA and GQA models side by side, with a memory estimator.
    What to read
    Read the Introduction and GQA Memory Savings sections. Run memory_estimator_gqa.py with two or three n_kv_groups values, then compare gpt_with_kv_gqa.py with gpt_with_kv_mha.py to see where the K/V heads are repeated. If you have time, read the Introduction in the sibling 05_mla folder for multi-head latent attention.
    What to skip
    Skip plot_memory_estimates_gqa.py and the full 32,768-token generation runs. The estimator shows the same result in seconds.

  2. Free alternative

    FlashAttention: Making Attention I/O-Aware (opens in a new tab)

    Hugging Face (community blog) · Aayush Garg · 45 min · Free · checked 21 Oct 2026

    Why
    Interviewers like to check that you know FlashAttention is exact and wins by moving less data, not by doing less maths. This post builds that case from the GPU memory hierarchy up.
    What to read
    Read The GPU Memory Hierarchy (including compute-bound vs. memory-bound), Standard Attention and Its IO Cost, and What FlashAttention Does: Operator Fusion, Tiling, Recomputation and Online Softmax. Afterwards, redraw the tiling loop with a running max and running sum from memory.
    What to skip
    On a first pass, skip the Standard backward pass derivation and The Logsumexp Value L. The forward-pass idea is what gets asked.

  3. Go deeper

    FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (opens in a new tab)

    arXiv · Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré · 60 min · Free · checked 21 Oct 2026

    Why
    If you claim IO-awareness in an interview, it helps to have seen the actual IO-complexity argument and Algorithm 1 yourself.
    What to read
    Read 2.1 Hardware Performance, 2.2 Standard Attention Implementation, 3.1 An Efficient Attention Algorithm With Tiling and Recomputation (with Algorithm 1) and 3.2 Analysis: IO Complexity of FlashAttention.
    What to skip
    Skip 3.3 Block-Sparse FlashAttention, the section 4 benchmarks and the appendix backward-pass details.

  4. Go deeper

    GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (opens in a new tab)

    arXiv · Joshua Ainslie et al. · 20 min · Free · checked 22 Oct 2026

    Why
    Grouped-query attention is in most current open models, and the interview question is usually why it’s a better trade-off than both multi-head and multi-query attention.
    What to read
    The paper is short; read all of §2 with Figures 1 and 2 (uptraining by mean-pooling heads, then grouped KV heads), §3.2 with Figure 3 (quality against inference time) and §3.3’s ablations on the number of groups and uptraining compute.
    What to skip
    Skip §3.1 setup detail and §4 related work.

Companies that test Transformers most

Put Attention variants in your plan

Tell us your interview date and hours a day; it lands on the right day, in the right order.

Build my plan