Chapter 1 LLM Handbook Part 1: Architecture, Attention Heads & Flash Attention
An engineer's guide to the core mechanics of Large Language Models. Demystifying the Transformer architecture, KV Cache, and how Flash Attention changed the game.
The Transformer Foundation
To optimize an LLM in production, you must understand what it is actually doing. At its core, an LLM (based on the Transformer architecture introduced in 2017) is an autoregressive engine: it predicts the next token based on all previous tokens.
It does this through stacked Layers, primarily composed of two operations:
- Self-Attention: Understanding context (how words relate to each other).
- Feed-Forward Networks (FFN): Recalling knowledge and facts.
In standard architectures (like LLaMA), a 70B model simply has more layers, wider hidden dimensions, and more attention heads than an 8B model.
Unpacking Attention Heads
When you read a sentence, you subconsciously link verbs to their subjects and pronouns to their nouns. “Attention” is how the model does this mathematically.
Instead of one monolithic attention mechanism, models use Multi-Head Attention (MHA). The model splits its focus into multiple “heads.” For example, a 70B model might have 64 attention heads per layer.
- Head 1 might focus on grammar.
- Head 2 might focus on subject-verb agreement.
- Head 64 might focus on emotional tone.
During inference, these heads generate three vectors for every token: Query (Q), Key (K), and Value (V).
- Query: What am I looking for?
- Key: What do I contain?
- Value: What is my actual semantic meaning?
The model takes the current token’s Query and does a dot-product multiplication against every previous token’s Key. High scores mean high relevance. The model then uses those scores to sum up the Values.
The Infrastructure Nightmare: The KV Cache
If the model had to recalculate the Keys and Values for every single word in a 10,000-word prompt just to generate the 10,001st word, inference would be unusably slow.
To solve this, engines like llama.cpp or vLLM use a KV Cache. Once a token’s Key and Value are calculated, they are saved in VRAM. When generating the next token, the model only calculates the Q, K, V for the new token, and pulls the past K and V from the cache.
The catch: The KV Cache grows linearly with context size. For a 256K context window, the KV cache can easily consume 10GB+ of VRAM. This is why aggressive quantization of the V-cache (to Q4) is critical for high-context local deployment.
Flash Attention: The Savior of Context
As context grows, the matrix multiplication required for attention (Q * K) grows quadratically ($O(N^2)$). For a 128K context, the math becomes paralyzing.
Flash Attention (and Flash Attention 2/3) solved this. It is a hardware-aware algorithm (originally designed for CUDA, now ported to SYCL/Metal) that radically minimizes memory reads/writes.
Instead of computing the massive Q*K matrix and writing it to High Bandwidth Memory (HBM), Flash Attention computes it in small “tiles” directly inside the GPU’s ultra-fast SRAM (L1 cache), calculates the Softmax, and only writes the final result back to HBM.
This turns a memory-bandwidth-bound operation into a compute-bound operation, enabling the massive 128K and 256K context windows we see today without slowing inference to a crawl.