Transformer Complexity, Inference Optimization, GQA and FlashAttention

In a fundamentals round for a staff-level AI engineering role, the interviewer asked a series of core Transformer questions: the computational complexity of a Transformer, how to optimize inference, and two specific techniques, Grouped-Query Attention and FlashAttention. Answer each part as you would out loud, with the key formulas and trade-offs. Notation you may use: sequence length $n$, model width $d$, $h$ attention heads of size $d_h = d / h$, $L$ layers, batch size $b$. ### Clarifying Questions - Should complexity be counted per layer or for the whole model, and for training, for processing a prompt, or for generating one token at a time? - Is the inference setting a decoder-only language model generating text online, or an encoder that scores inputs in batches? - What matters most for inference here: latency per request, throughput per GPU, memory, or cost? ### Part 1 — Complexity of a Transformer Give the time and memory complexity of one Transformer layer as a function of $n$ and $d$, say which term dominates for short and for long sequences, and explain how the cost changes during autoregressive generation. ```hint Count the matrix multiplications Separate the work that scales with the number of tokens times the width squared from the work that compares every token with every other token. ``` #### What This Part Should Cover - The FLOP count of the projections, the attention scores and weighted sum, and the feed-forward block - Memory for parameters, activations and the attention matrix, and where the crossover between the linear and quadratic terms lies - The cost per generated token with and without cached keys and values ### Part 2 — Inference optimization How would you make Transformer inference faster and cheaper? ```hint Find the bottleneck first Processing a whole prompt and generating one token at a time stress the hardware differently. Decide which resource each phase runs out of before you pick techniques. ``` #### What This Part Should Cover - The difference between the prompt (prefill) phase and token-by-token decoding, and the bottleneck of each - Techniques at the model, memory, kernel and serving levels, each tied to the bottleneck it relieves - The quality, latency and cost trade-off of each technique ### Part 3 — Grouped-Query Attention Explain Grouped-Query Attention: how it differs from standard multi-head attention and from multi-query attention, what it saves, and what it costs. ```hint Look at what is kept in memory During token-by-token decoding, think about which tensors are stored for every past token, and how their size grows with the number of heads. ``` #### What This Part Should Cover - The mechanism, and the two extremes it interpolates between - The effect on KV-cache memory, memory bandwidth per decoding step, parameters and FLOPs - The quality trade-off, and how an existing multi-head model can be converted ### Part 4 — FlashAttention Explain FlashAttention: what problem it solves, how it computes attention faster, and how its compute and memory compare with a standard implementation. ```hint Count memory traffic, not FLOPs Ask where the $n \times n$ score matrix lives in a standard implementation, and how many times it moves between GPU high-bandwidth memory and on-chip memory. ``` #### What This Part Should Cover - Why standard attention is limited by memory traffic rather than arithmetic - How the computation is reorganized so the full score matrix never reaches high-bandwidth memory, in both the forward and backward passes - Its FLOPs and memory compared with standard attention, whether the result is exact, and where it helps less ### What a Strong Answer Covers - Correct complexity with the constants that matter, separated into training, prefill and decoding - Optimizations chosen from the identified bottleneck rather than recited as a list - Precise mechanisms for Grouped-Query Attention and FlashAttention, including what each does not change - Quantified trade-offs: memory saved, quality risk and engineering cost - How the techniques combine in one serving stack ### Follow-up Questions - At long context lengths, when does the KV cache rather than the weights dominate GPU memory during decoding, and how would you estimate the crossover? - Does FlashAttention change the asymptotic complexity of attention? What would you use if you needed sub-quadratic cost? - How would you choose the number of key/value heads for a new model? - Why does batching help decoding throughput much more than it helps prefill?