What self-attention actually computes, why it is memory-bandwidth-bound not compute-bound, and why that determines which hardware wins the inference era
In Lesson 6 we learned the H100's HBM3 bandwidth is 3.35 TB/s, but its Tensor Core throughput
is 3,958 TFLOPS. A chip is compute-bound when computation fills time; it is
memory-bandwidth-bound when waiting for data to arrive from DRAM fills time.
The crucial fact about Transformers at inference time: they are almost entirely
memory-bandwidth-bound. Understanding why tells you what matters in inference
hardware (bandwidth and HBM capacity, not raw TFLOPS) β and where the next wave of
semiconductor investment is going.
Part 1 β Why Transformers Exist: The Problem They Solve
RNNs couldn't see long-range context. Attention can β but at a cost.
Before 2017, sequence modeling meant recurrent neural networks (RNNs/LSTMs). An RNN
processes tokens one by one, left to right, compressing everything seen so far into a
fixed-size hidden state vector. The pathology: information from position 1 has to travel
through hundreds of RNN cells to reach position 500. By then it has been overwritten by
everything in between. RNNs could not learn long-range dependencies reliably β the gradient
signal that taught the network about position 1 had to backpropagate through 500 steps,
diminishing to near-zero (vanishing gradient). LSTMs helped but didn't solve it structurally.
The 2017 "Attention Is All You Need" paper from Google Brain proposed a radical alternative:
throw away recurrence entirely. Instead, let every token look at every other token
simultaneously, with a learned weight for how much attention to pay each one.
This is the self-attention mechanism. The price: instead of O(n) memory (one hidden state),
you need O(nΒ²) memory β every pair of tokens. For n=8192 tokens, that is 67 million pairs.
This trade-off β parallelism and long-range context for quadratic cost β defines
everything about Transformer hardware requirements.
Part 2 β The Attention Equation: Every Token Asking a Question
Attention(Q, K, V) = softmax(QKα΅ / βd_k) Β· V
This one equation is the entire self-attention mechanism. Let's unpack every term from scratch.
Step 1: Embedding β tokens become vectors
Every input token (a word-piece, roughly 3β4 characters) is converted into a vector of
d_model numbers β
the embedding. For Llama 3 8B: d_model = 4,096. For GPT-4 class: d_model β 12,288.
This vector is the token's identity β learned during training to encode semantic meaning.
"King" and "Queen" will have embeddings that are geometrically close and differ mainly in
a "royalty dimension."
d_model
Model dimension
The length of every token's vector. Llama 3 8B: 4,096. GPT-4 class: ~12,288.
H
Attention heads
Number of attention heads. Llama 3 8B: 32 heads. Each head learns different relationships.
d_k
Head dimension
d_model / H. Llama 3 8B: 4,096 / 32 = 128. Each head works in this smaller space.
n
Context length
Number of tokens in sequence. Llama 3: 8K. GPT-4: 128K. Claude 3: 200K. The nΒ² problem.
Step 2: Q, K, V projections β what are they physically?
For each token, the embedding is linearly projected three ways using learned weight matrices:
Q / K / V are linear projections of the same input embedding x
Q = x Β· W_Q β "What am I looking for?" (query)
K = x Β· W_K β "What information do I contain?" (key)
V = x Β· W_V β "What do I actually deliver if selected?" (value)
All three are the same input x; different learned weight matrices W_Q, W_K, W_V produce different projections.
Shape of each: (n, d_k) β one d_k-dimensional vector per token, per head.
The database analogy holds precisely: Q is the query you submit; K is the index entry of
each row; V is the row's content. Attention does a "soft database lookup" β instead of
returning one row exactly, it blends all rows weighted by how well each K matches the Q.
Crucially, the matching is learned: the network learns what Q/K/V
projections are useful for predicting the next token.
Step 3: The full attention equation
Self-attention (one head) β step by step
1.S = Q Β· Kα΅β³ Dot product of every Q with every K. Shape: (n, n). The nΒ² matrix.
2.S = S / βd_kβ³ Scale by βd_k to prevent dot products from growing too large (softmax saturation).
3.A = softmax(S, dim=-1)
β³ Convert scores to probabilities. Each row sums to 1.0. Shape still (n, n).
4.Output = A Β· Vβ³ Weighted sum of Values. Shape: (n, d_k). The context-aware representation.
Attention(Q,K,V) = softmax(QKα΅ / βd_k) Β· V
The most important thing to understand about this equation
Position 500 in a sequence can directly attend to position 1 β with a single matrix multiply,
no information loss, no gradient attenuation across 499 steps. The score S[500, 1] captures
exactly how relevant token 1 is to token 500. This is why Transformers learned
long-range dependencies that RNNs couldn't β but it required storing an (n Γ n) matrix.
At n = 128,000 tokens (GPT-4 context), S is a 128K Γ 128K matrix = 65 billion attention
scores. At BF16 (2 bytes), that is 131 GB β larger than the H100's entire 80 GB HBM.
This is the memory wall that FlashAttention and KV caching both exist to address.
Interactive: Simulated attention weights
Select a token below. The heatmap shows which tokens it attends to most strongly.
This is a simplified illustration of real attention patterns learned by GPT-2 on English text.
Attention weights for selected token β (darker = higher weight)
Click a token above to see its attention pattern.
Part 3 β Multi-Head Attention: Parallel Relationship Detectors
H independent attention heads, each learning a different type of relationship
A single attention head can only capture one type of relationship per token pair.
Multi-head attention runs H heads in parallel with different W_Q, W_K, W_V weight matrices β
each head learns to attend to a different kind of pattern. After computing output for each
head, results are concatenated and projected back to d_model.
32 heads run in parallel on the GPU (batched as a 3D tensor operation). Each head has its own
W_Q, W_K, W_V, so each learns different relationships. Research has found individual attention
heads specialize: some track syntactic dependencies (subjectβverb), some track coreference
("she" β "the queen"), some track positional patterns (each token attends to the previous one).
No one programs this β it emerges from training.
A single Transformer layer = one MHA + one FFN (feed-forward network), each
wrapped in a residual connection and layer normalization. Llama 3 8B has 32 such layers
stacked. GPT-4 class: ~96 layers. Each layer runs the full attention over all tokens.
Pre-norm Transformer block (modern architecture used in Llama 2/3, Mistral, Gemma)
Two important architecture notes: (1) Pre-norm (LayerNorm before each sub-layer)
is what modern models use β it trains more stably than the original "post-norm" in the 2017 paper.
(2) The FFN block β two large linear transformations with a nonlinearity β is where
most of the model's parameters live. In Llama 3 8B, the FFN intermediate dimension is 14,336 (not
4Γ4096 = 16,384 as in original Transformers, due to SwiGLU gating). The FFN is essentially the
model's "memory" β MHA routes information, FFN stores and recalls patterns.
Part 5 β Why Inference Is Memory-Bandwidth-Bound (Not Compute-Bound)
Arithmetic intensity: the ratio that determines what's the bottleneck
This is the most important section for connecting the Transformer to hardware investment.
Arithmetic intensity = floating point operations / bytes of memory accessed.
If a GPU's ratio (peak FLOPS / memory bandwidth) exceeds a workload's arithmetic intensity,
the workload is memory-bandwidth-bound. The workload is waiting for data to arrive from HBM,
not for computation to complete. Adding more Tensor Cores doesn't help β you need more
memory bandwidth or higher cache hit rate.
A workload with arithmetic intensity > 1,182 ops/byte β compute-bound. Below 1,182 β memory-bandwidth-bound.
Attention at inference β arithmetic intensity calculation
Decoding one token (batch=1): load all model weights + KV cache from HBM, do one pass.
Llama 3 8B weights: 8B params Γ 2 bytes (BF16) = 16 GB to load each decode step
Compute per step: ~8B ops (one matmul per weight) β 8 Γ 10βΉ ops / 16 Γ 10βΉ bytes
Arithmetic intensity at batch=1: ~0.5 ops/byte β 2,000Γ below the H100's ridge point of 1,182 ops/byte
At batch=1, the H100 is memory-bandwidth-bound. Tensor Cores sit idle ~99.95% of the time. You need bandwidth, not TFLOPS.
How batch size changes the equation
Batch=1 (chatbot, single user)
0.5
0.5 ops/B
Batch=16 (small API cluster)
8
8 ops/B
Batch=256 (busy inference API)
128
128 ops/B
Batch=2,048 (datacenter throughput)
1,024
1,024 ops/B
H100 ridge point (compute=BW limited)
1,182 β reach here and Tensor Cores become the bottleneck
1,182 ops/B
Inference serving is almost always batch=1 to batch=64 (latency-sensitive APIs). Most inference workloads sit in the red/orange zone β deeply memory-bandwidth-bound.
Training typically runs batch β₯ 1,024 β closer to compute-bound.
This is why different hardware designs dominate training vs. inference.
The investment implication: why inference hardware is different from training hardware
Training: high batch sizes β arithmetic intensity near the ridge point β Tensor Core TFLOPS matter β
NVIDIA H100/B200 with 3,958 TFLOPS wins.
Inference: low batch sizes β deeply memory-bandwidth-bound β bandwidth and HBM capacity matter more
than TFLOPS β opens the door for specialized inference chips.
Groq's LPU (Language Processing Unit) is designed entirely around this insight: it has ~80 TB/s
on-chip SRAM bandwidth (vs. H100's 3.35 TB/s HBM + 20 TB/s shared memory), enabling ~10Γ lower
latency for single-user inference at the cost of much smaller model capacity.
AMD's MI300X (192GB HBM3) wins for inference of large models that don't fit in 80GB H100.
NVIDIA's response: H200 (141GB HBM3e, 4.8 TB/s), B200 (192GB HBM3e, 8 TB/s).
The inference era is a memory arms race, not a TFLOPS race.
Part 6 β The KV Cache: Why Long Context Explodes Memory
The mechanism that makes serving 200K-context models extremely expensive
During inference (generating text token by token), a Transformer is autoregressive:
each new token is generated by running attention over all previous tokens.
Without optimization, this means re-computing K and V for all previous tokens at every step β
absurdly wasteful since they haven't changed.
The KV cache caches the Key and Value matrices for all previous tokens,
reusing them for each new token generated. Only Q changes per step. This reduces computation
from O(nΒ²) per token to O(n) per token β a massive win for speed.
The cost: you must store the growing KV cache in HBM as generation proceeds.
Model weights alone: 8B Γ 2 bytes = 16 GB (already fills most of an 80GB GPU) For 128K context: weights (16GB) + KV cache (17.2GB) = 33.2 GB per concurrent request How many concurrent 128K-context requests on one 80GB H100? β 80 / 33.2 β 2 requests
For GPT-4 class (96 layers, 128 KV heads, 128K ctx, d_head=128): KV cache β 150 GB per request. Requires distributing across multiple H100s per user.
This is why Claude's 200K context window and GPT-4's 128K context window are so
computationally expensive to serve β not the inference TFLOPS but the HBM required
to store the KV cache for each concurrent user. More HBM capacity = more concurrent
users = more revenue per GPU.
Grouped Query Attention (GQA): the KV cache solution
Multi-Head Attention (MHA)
Original Transformer (2017), GPT-2, BERT. H Q heads, H K heads, H V heads.
KV heads = H = 32
KV cache is largest. Every head has its own K and V stored. High expressiveness, high memory cost.
Used by: GPT-2, BERT, Llama 1
Grouped Query Attention (GQA)
G groups of Q heads share one K/V pair each. If H=32 Q heads and G=4 groups: 8 K/V pairs serve 32 Q heads (4 Q heads per K/V).
KV heads = H/G = 32/4 = 8
KV cache is 4Γ smaller. Minimal quality loss. The modern default.
Extreme case of GQA: G=H, so 1 K/V pair shared by all H Q heads. Maximum KV cache reduction.
KV heads = 1
KV cache is 32Γ smaller than MHA. Some quality degradation on complex tasks. Very fast inference.
Used by: PaLM, Falcon, some Mistral variants
The trend is clear: every major model released in 2024β2025 uses GQA or MQA.
This is not an architectural choice for quality β it is an engineering choice for
inference cost. GQA reduces KV cache size by 4β32Γ, allowing 4β32Γ more concurrent
users per GPU, directly improving inference economics.
Part 7 β FlashAttention: Rewriting the Attention Algorithm for Hardware Reality
The most important AI systems paper of the last 5 years
The naive attention implementation materializes the full (nΓn) attention score matrix S in HBM.
At n=8,192 tokens, S is 8192Β² Γ 2 bytes = 134 MB per head per layer. For Llama 3 8B:
32 heads Γ 32 layers = 1,024 attention matrices = 137 GB of HBM reads/writes just for the
attention scores. At 3.35 TB/s HBM bandwidth, this takes ~41ms per forward pass β
before any actual computation.
FlashAttention (Dao et al., 2022; FA2 2023; FA3 2024) makes the key observation:
you never actually need to store the full S matrix. You can compute the final output
in tiles that fit in SRAM (shared memory on the SM), accumulating the softmax
normalization incrementally. The online softmax algorithm enables this.
Standard attention vs. FlashAttention β memory access pattern
FlashAttention doesn't change the mathematical result β it changes the order of operations to
avoid writing the intermediate S matrix to HBM. The key insight: softmax can be computed
incrementally using an "online normalization" technique (running max + sum), so tiles of
QΓKα΅ can be processed one at a time in SRAM, accumulating the weighted V sum without
ever materializing the full nΓn attention matrix. FlashAttention 3 (2024) achieves 75% of
H100 HBM bandwidth utilization β the highest ever achieved for an attention kernel.
This is now built into every major ML framework (PyTorch via scaled_dot_product_attention,
JAX, and all production LLM serving frameworks).
Part 8 β Position Encoding: Teaching the Transformer Where Things Are
Self-attention is permutation-invariant. You have to tell it about order.
The attention matrix computes a dot product between every Q and every K β but the
dot product is the same regardless of order. If you shuffle "the dog chased the cat"
into "cat the the dog chased," the attention scores are identical (same tokens, same QΓK).
Position information must be added explicitly.
The original 2017 paper used sinusoidal encoding β fixed patterns added to the embedding.
Modern models use RoPE (Rotary Position Embedding), introduced in the
RoFormer paper (2021) and now used in Llama, Mistral, Falcon, and almost all
new models. RoPE encodes position by rotating the Q and K vectors in 2D planes before
computing the dot product, such that the dot product naturally depends on the
relative position (token i vs. token j), not absolute position.
This makes RoPE much better at extrapolating to contexts longer than the model was
trained on β which is why all 128K+ context models use RoPE variants (YaRN, LongRoPE).
Part 9 β Investment Implications of Inference Architecture
The inference era reshapes the hardware stack β and the investable landscape
Training is a one-time event per model. Inference is continuous. As models transition from
"being trained" to "being deployed," the economic center of gravity shifts from training hardware
to inference hardware. McKinsey estimates inference will represent ~60β70% of AI compute spend
by 2027. The memory-bandwidth-bound nature of inference creates a different winner's landscape.
~60%
Inference share of AI compute by 2027E
Up from ~30% in 2023. The training spike was the capex wave; inference is the recurring subscription.
8 TB/s
B200 HBM3e bandwidth
2.4Γ H100's 3.35 TB/s. NVIDIA's answer to the inference memory arms race.
192 GB
B200 HBM3e capacity
2.4Γ H100's 80GB. Allows 2.4Γ more concurrent users or 2.4Γ longer context per user.
SK Hynix
Primary HBM supplier
~50% of HBM market. B200's 8 TB/s requires HBM3e β SK Hynix ships first.
The inference hardware landscape β who wins what
NVIDIA H200 / B200 β the default inference chip
β Massive HBM capacity (141GB / 192GB) for large model serving
Amazon Inferentia 3 / Google TPU v5e β hyperscaler inference
β Custom-optimized for hyperscaler's own model shapes
β 3β5Γ better performance/watt for specific workloads
β Dramatically reduces GPU spend for internal workloads
β Not available on open market (no investable exposure)
β Requires model compilation step β not plug-and-play
Wins: Internal hyperscaler inference. Reduces NVIDIA's TAM at the margin.
The compound thesis: HBM is the scarcest resource in the inference era
Every trend in Transformer architecture points toward needing more HBM:
larger models (more weights to store), longer contexts (KV cache grows linearly),
more concurrent users (more KV caches in parallel), higher quality (bigger models).
The inference era is fundamentally a memory capacity and bandwidth arms race β not a TFLOPS race.
The supply chain bottleneck: HBM is made by three vendors (SK Hynix ~50%, Samsung ~35%, Micron ~15%).
HBM requires 12-high stacking with TSVs (Lesson 5 β CoWoS technology) and takes 6+ months to ramp.
If the inference era grows as projected, SK Hynix's HBM ASP expansion and volume growth
is arguably the most direct play on LLM inference economics β arguably more directly
than NVIDIA itself, since HBM is needed regardless of which GPU vendor wins.
SK Hynix (000660.KS) reported HBM revenue growing 300%+ YoY in 2024 and is the primary
HBM3E supplier for NVIDIA B200.
Practice Project β Attention Calculator
Given a model and deployment scenario, calculate the KV cache memory requirement and
determine the inference mode (memory-bandwidth-bound vs. compute-bound).
Technical quiz
1. The self-attention equation is Attention(Q,K,V) = softmax(QKα΅ / βd_k) Β· V. What does the βd_k scaling factor prevent?
It normalizes the output so it has unit variance, matching layer normalization's output range
It prevents dot products from growing proportionally to d_k (since they sum d_k random variables), which would push softmax into saturation where gradients vanish
It converts raw attention scores into probabilities so they can be multiplied by V
It reduces memory use by shrinking the intermediate attention matrix before the softmax
2. A Transformer with n=4,096 tokens, 32 layers, and 32 attention heads needs to materialize the full attention score matrix S in HBM using naive (non-FlashAttention) implementation. How much HBM does S require in BF16?
3. Why does serving a GPT-4 class model at 128K context length require distributing a single user's request across multiple H100 GPUs (tensor parallelism), even ignoring compute time?
128K tokens require more TFLOPS than one H100 can deliver within latency SLAs
The KV cache for one GPT-4 class request at 128K context exceeds 150 GB β larger than one H100's 80 GB HBM β so the KV cache must be sharded across multiple GPUs
FlashAttention cannot tile attention for sequences longer than 32K tokens on one GPU
The model's parameters (weights) alone exceed 80 GB and must be sharded
4. Grouped Query Attention (GQA) reduces the KV cache size by sharing K and V heads across multiple Q heads. What is the primary trade-off compared to standard Multi-Head Attention (MHA)?
GQA requires more compute per token because Q heads must project into a lower-dimensional space
GQA cannot be used with FlashAttention because tiling requires equal Q and KV head counts
GQA reduces the expressiveness of attention β fewer K/V heads means less capacity to learn distinct relationship patterns β with empirically small quality loss at scale but some degradation on complex tasks
GQA requires a more complex training procedure because shared KV matrices have conflicting gradient signals from multiple Q heads
5. From an investment perspective, why is SK Hynix's HBM business arguably more structurally attractive than NVIDIA in the inference era?
SK Hynix has a higher gross margin than NVIDIA, making it fundamentally more profitable per dollar of revenue
SK Hynix serves all GPU vendors (NVIDIA, AMD, Intel, Google TPU) with HBM, giving it broader customer exposure and less concentration risk
HBM demand is non-discretionary for inference (every GPU needs HBM regardless of vendor), SK Hynix has ~50% market share in a 3-vendor oligopoly, and inference is fundamentally memory-bandwidth-bound β meaning HBM is the constrained resource, not Tensor Core TFLOPS
SK Hynix is trading at a lower P/E than NVIDIA, making it the better value even if the growth rates were identical
Primary sources
"Attention Is All You Need" (Vaswani et al., 2017) β The original Transformer paper. Read Section 3 (model architecture) and Section 3.2.1 (scaled dot-product attention). Dense but completely self-contained.
FlashAttention (Dao et al., 2022) β The paper that made long-context Transformers practical. The introduction and Section 2 (background) explain the memory hierarchy argument beautifully.
Lesson 8: LLM inference at scale β vLLM, PagedAttention, continuous batching, speculative decoding, and how inference serving systems squeeze 10Γ more throughput from the same hardware
Ask me anything. Good follow-ups for a tech worker:
"Walk me through implementing FlashAttention's online softmax algorithm" Β·
"Why does RoPE generalize to longer contexts better than sinusoidal encoding?" Β·
"What is tensor parallelism and pipeline parallelism, and when do you use each?" Β·
For an investor:
"Model SK Hynix's HBM revenue growth if inference compute doubles each year" Β·
"How does speculative decoding change the hardware requirements for inference?" Β·
"What is the right valuation framework for an inference-focused chipmaker vs. a training-focused one?"