Transformer Architecture Overview

Every optimization in AI infrastructure acts on some component of the Transformer. Before diving into the details, this chapter provides the map — what are the building blocks, how do they fit together, and why does each one matter for infrastructure.

---

1. Why AI Infra Engineers Must Understand the Transformer

The Transformer is not just a model architecture — it’s the compute workload that defines all AI infrastructure requirements. Every infrastructure decision maps to a Transformer component:

Infra layer	Transformer component	Challenge
CUDA kernels	Attention, FFN, Softmax, LayerNorm	Optimizing individual operations (FlashAttention, fused kernels)
Distributed training	Attention heads, FFN weights, layers	Splitting the model across GPUs (tensor/pipeline/expert parallelism)
Inference serving	KV cache, autoregressive decode	Managing memory, batching, latency optimization
Hardware design	Matrix multiplications, memory access	Tensor Cores, HBM bandwidth, NVLink for parallelism

Understanding which component dominates which cost tells you where to optimize.

---

2. The Building Blocks

A modern decoder-only LLM has surprisingly few distinct components:

Input text
  → Tokenizer (text → token IDs)
  → Embedding table (token IDs → vectors)
  → N × Decoder Block:
      → RMSNorm → Masked Self-Attention → Residual
      → RMSNorm → SwiGLU FFN → Residual
  → Final RMSNorm
  → LM Head (vectors → vocabulary logits)
  → Sampling (logits → next token)

Only 7 distinct operation types: tokenization, embedding lookup, normalization, attention, FFN, linear projection, and softmax/sampling. Everything else is just these operations repeated and composed.

Component Map

Component	What it does	Parameters	Compute cost	Memory concern
Embedding	Token ID → vector	$V \times d$	Negligible (lookup)	Large if vocab is big
Self-Attention	Token-to-token interaction	$4d^2$ per layer	$O(S^2 d)$	KV cache grows with sequence length
FFN (SwiGLU)	Per-token processing	$3d \times d_{ff}$ per layer	$O(S \cdot d \cdot d_{ff})$	Activation memory
RMSNorm	Stabilize activations	$d$ per instance	Negligible	—
LM Head	Vector → vocab probabilities	$d \times V$	One GEMM	Often tied with embedding

Where the parameters live: In a typical 7B model, ~66% of parameters are in FFN weights, ~33% in attention weights, and ~1% in everything else. In a 70B model with GQA, FFN dominance increases further.

Where the compute goes: For short sequences, FFN dominates compute. For long sequences, attention’s $O(S^2)$ term takes over. The crossover point depends on model dimensions and sequence length.

---

3. How It Evolved

The Three Architecture Families

The original Transformer (2017) had an encoder-decoder design for machine translation. Three families emerged:

Architecture	Key models	How it works	Status
Encoder-only	BERT, RoBERTa	Bidirectional attention, masked token prediction	Superseded for generation
Encoder-Decoder	T5, BART, Whisper	Encoder processes input, decoder generates	Used for translation, ASR
Decoder-only	GPT, LLaMA, Qwen, Mistral	Causal attention, next-token prediction	Dominant for LLMs

Why Decoder-Only Won

1. Simplicity: One stack of identical blocks — no encoder, no cross-attention
2. Scalability: Just add more blocks. Same architecture from 1B to 400B+
3. Unified: Next-token prediction handles every task when prompted correctly
4. Inference efficiency: KV cache works naturally; encoder-decoder needs two separate caches

Key Architecture Innovations Since 2017

Innovation	What changed	Why it matters for infra
GQA (2023)	Fewer KV heads than Q heads	4-8× KV cache reduction
RoPE (2021)	Rotary position embedding	Enables long-context extension
SwiGLU (2020)	Gated FFN activation	3 weight matrices instead of 2
RMSNorm	Simpler normalization	10-15% faster than LayerNorm
MoE (2022+)	Sparse expert routing	10× parameters at ~1× compute
MLA (2024)	Latent KV compression	Further KV cache reduction
MTP (2024)	Multi-token prediction heads	Enables speculative verification

---

4. Reading Order for This Section

The remaining articles in this section build on each other:

1. Tokenization — How text becomes numbers (the input)
2. Transformer as Compute Graph — Full forward pass with shapes and FLOPs
3. Self-Attention — The core $O(S^2)$ operation
4. FFN & Normalization — The other half of each block
5. Positional Encoding — How the model knows token order (RoPE)
6. Autoregressive Generation — How the model actually generates text (prefill, decode, KV cache)
7. Attention Variants — GQA, MLA, Mamba — reducing KV cache cost
8. MoE Architecture — Sparse models for more capacity at less compute

Each article includes shape tracking, compute analysis, and connections to infrastructure challenges.

---

Self-Check

• Name the 7 distinct operation types in a decoder-only Transformer
• Explain where parameters and compute concentrate (FFN vs attention)
• Explain why decoder-only architecture became dominant
• Name 3 architecture innovations since the original Transformer and their infrastructure impact
• Explain the crossover between FFN-dominated and attention-dominated compute as sequence length grows

---

References

• Attention Is All You Need (Vaswani et al., 2017)
• Language Models are Few-Shot Learners (GPT-3, Brown et al., 2020)
• LLaMA (Touvron et al., 2023)