Large Language Models University

Module 2.1: Word Coordinates — High-Dimensional Embeddings

In geometry class, you plot points on a 2D graph using $(x, y)$ coordinates like $(3, 5)$. If you want to describe a point in 3D space (like a drone flying in a room), you add a third coordinate $z$ for height: $(x, y, z)$.

What if you wanted to describe the meaning of a word mathematically? In an LLM, we give every word a vector with hundreds or thousands of coordinates (typically 4,096 dimensions)! This list of numbers is called an Embedding.

In this high-dimensional semantic space, words that share similar meanings live close together. Words like "puppy" and "dog" have coordinates that are nearly identical, while "refrigerator" is far away!

To calculate how closely related two words are, we compute the Cosine Similarity between their vectors. If the angle between them is $0^\circ$, their cosine similarity is $1.0$ (perfect match). If they are completely unrelated ($90^\circ$ perpendicular), the similarity is $0.0$!

Module 2.2: The Attention Spotlight — Resolving Context

Consider this sentence: "The animal didn't cross the street because it was too tired."

What does the word "it" refer to? As a human reader, you know instantly that "it" refers to the animal, because streets don't get tired! But what if the sentence ended with "...because it was too wide"? Suddenly, "it" refers to the street!

Traditional computers would get hopelessly confused. But the Transformer architecture uses a mechanism called Self-Attention. When processing the word "it", the model shines an attention spotlight over every single previous word, checking which words provide the necessary clues to clarify its meaning.

Module 2.3: Temperature & The Softmax Probability Distribution

Once the LLM scores all possible next tokens, it has a list of raw numerical points called Logits. To convert these logits into clean percentages that add up to $100\%$, it applies the Softmax function:

The variable $\tau$ (tau) is the Temperature:

• Low Temperature ($\tau = 0.1 - 0.3$): Makes the top-ranked word dominant ($>95\%$). The output is predictable, factual, and strictly logical. Great for coding and math!
• Medium Temperature ($\tau = 0.7 - 0.8$): Balances coherence with natural variety. The standard setting for conversational chatbots.
• High Temperature ($\tau \ge 1.2$): Gives rare and unusual words a higher chance of being chosen. Highly creative, but with a high risk of wild tangents or hallucinations!

Module 3.1: Scaled Dot-Product & Multi-Head Attention (MHA)

In 2017, researchers published the landmark paper "Attention Is All You Need", replacing recurrence and convolutions with the Transformer. At its heart lies Scaled Dot-Product Attention using three linear projections: Queries ($Q$), Keys ($K$), and Values ($V$):

Why divide by $\sqrt{d_k}$? As the key dimension $d_k$ grows large, the dot products grow proportionally in magnitude. Large dot products push the Softmax function into saturation regions with microscopic gradients ($\approx 0$), causing the vanishing gradient problem during backpropagation. Dividing by $\sqrt{d_k}$ stabilizes the variance to exactly $1.0$!

In Multi-Head Attention (MHA), rather than computing attention once across the full hidden dimension $d_{\text{model}}$, the vectors are split into $h$ parallel "heads" of dimension $d_k = d_{\text{model}} / h$. This allows head 1 to focus on grammar, head 2 on entities, and head 3 on pronouns simultaneously!

Module 3.2: Positional Encodings (RoPE) & SwiGLU MLPs

Unlike Recurrent Neural Networks (RNNs) that read text sequentially word-by-word, a Transformer processes all tokens simultaneously in parallel. But language has syntax: "Dog bites man" is very different from "Man bites dog"! Without positional information, attention is completely permutation-invariant.

Modern LLMs (such as LLaMA, Mistral, and Gemma) use Rotary Position Embeddings (RoPE). Instead of adding fixed position vectors, RoPE rotates the Query and Key vectors in 2D coordinate pairs by an angle proportional to the token position $m$:

This ensures that the inner product $(R_m Q) \cdot (R_n K)^T$ depends strictly on the relative distance $(m - n)$ between tokens rather than their absolute index.

Following attention, tokens pass through a Feed-Forward Network (FFN). Modern LLMs use the SwiGLU activation function:

Module 3.3: Pre-training vs Fine-Tuning & Alignment (RLHF / DPO)

Building a production-ready LLM requires three distinct phases:

1. Pre-training: Self-supervised training on 5 to 15 trillion tokens using Cross-Entropy Loss: $$\mathcal{L}_{\text{pretrain}} = -\frac{1}{T} \sum_{t=1}^T \log P(x_t \mid x_{
2. Supervised Fine-Tuning (SFT): Training on high-quality question-and-answer pairs so the base model behaves like an instruction-following assistant.
3. Alignment (RLHF & DPO): Aligning responses with human preferences. Direct Preference Optimization (DPO) eliminates unstable reinforcement learning reward models by directly optimizing the policy using implicit probabilities between preferred ($y_w$) and dispreferred ($y_l$) completions.

Module 4.1: Prefill vs. Decode & The Roofline Bottleneck

Running LLM inference is not a monolithic workload. It consists of two fundamentally different phases with opposite hardware bottlenecks:

1. Prefill (Prompt Processing): The model takes in the user's entire prompt of length $S$ at once. The operations are dense Matrix-Matrix multiplications (GEMM). The arithmetic intensity is high ($\gg 100 \text{ FLOPs/Byte}$), meaning the GPU compute units (Tensor Cores) are fully saturated. Prefill is Compute-Bound!
2. Decode (Token Generation): The model generates tokens autoregressively one by one. For every single generated token, every single parameter in the entire model must be read from High Bandwidth Memory (HBM) just to perform a Vector-Matrix multiplication (GEMV)! Generating 1 token requires streaming ~140 GB of weights for a 70B parameter model in FP16, yielding an arithmetic intensity of only $\approx 1-2 \text{ FLOPs/Byte}$. Decode is Memory-Bandwidth Bound!

Under the Roofline Model, buying a GPU with $2\times$ faster ALUs will produce zero speedup during single-stream token generation unless HBM memory bandwidth ($BW_{\text{mem}}$) also increases!

Module 4.2: The Key-Value (KV) Cache Architecture

In standard attention, calculating $\text{softmax}(Q K^T) V$ requires the key and value vectors of all preceding tokens. Without caching, generating token $t$ would require recomputing the $K$ and $V$ vectors for all tokens $1 \dots t-1$, incurring quadratic $O(S^2)$ wasted compute!

To prevent this, serving engines allocate a Key-Value (KV) Cache in GPU VRAM, retaining the computed keys and values across every layer:

For a model like LLaMA-2 70B ($L=80$, $h=64$, $d_k=128$) with a context window of $8,192$ tokens and batch size $b=16$ in FP16:

$$\text{Memory}_{\text{KV}} = 2 \times 80 \times 64 \times 128 \times 8,192 \times 16 \times 2 \text{ bytes} \approx 343.6 \text{ GB!}$$

The KV cache alone exceeds the total VRAM of four 80GB GPUs!

Module 4.3: Multi-Query (MQA) & Grouped-Query Attention (GQA)

To break this memory wall, hardware and model architects developed alternative attention head topologies:

• Multi-Head Attention (MHA): Each Query head has its own dedicated Key and Value head ($n_{\text{kv}} = n_{\text{q}}$). Maximum representational expressiveness, but disastrous KV memory consumption.
• Multi-Query Attention (MQA): All Query heads share a single Key head and a single Value head ($n_{\text{kv}} = 1$). Slashes KV cache size by up to $64\times$, but causes measurable quality degradation in complex multi-step reasoning.
• Grouped-Query Attention (GQA): The ideal compromise adopted by LLaMA-3, Mistral, and DeepSeek. Query heads are partitioned into $G$ groups, with each group sharing 1 Key-Value pair ($n_{\text{kv}} = n_{\text{q}} / G$, typically $G=8$). Provides an $8\times$ reduction in KV cache memory and bandwidth traffic with virtually zero degradation in benchmark accuracy!

Module 5.1: FlashAttention-1, 2, and 3 — IO-Aware Tiling & Online Softmax

On modern GPUs, the memory hierarchy features a massive bandwidth gap: on-chip SRAM provides over $19 \text{ TB/s}$ of bandwidth (tens of megabytes capacity), while off-chip HBM provides only $\sim 2-3.35 \text{ TB/s}$ (tens of gigabytes capacity). Standard attention implementations materialize the full $S \times S$ attention matrix directly in slow HBM, incurring $O(S^2)$ memory reads and writes!

FlashAttention solves this by fusing the operations into a single GPU kernel. It tiles Query ($Q$), Key ($K$), and Value ($V$) into SRAM blocks ($B_r \times B_c$) and uses the Online Softmax algorithm to compute partial softmax statistics without ever storing the full intermediate attention matrix:

By keeping intermediate accumulations entirely within on-chip SRAM, FlashAttention reduces HBM memory traffic from $O(S^2)$ down to $O(S)$ linear complexity, achieving a $2\times - 4\times$ wall-clock speedup while using zero auxiliary memory!

FlashAttention-3 pushes this further on Hopper/Blackwell architectures by exploiting FP8 Tensor Core hardware pipelines, asynchronous Tensor Memory Accelerator (TMA) copies, and warp-specialization to overlap matrix multiplication with softmax exponentiation.

Module 5.2: PagedAttention & Continuous Dynamic Batching (vLLM)

In real-world LLM serving, traditional systems pre-allocated contiguous memory blocks for the maximum possible sequence length ($S_{\max} = 4,096$ or $8,192$). Because most requests terminate early, 60% to 80% of GPU memory was completely wasted due to internal and external memory fragmentation!

Inspired by classical operating system virtual memory, PagedAttention divides the KV cache into fixed-size physical blocks (e.g., 16 tokens per block). A software Block Table maps logical token indices to non-contiguous physical pages in VRAM:

Module 5.3: Post-Training Quantization (AWQ, GPTQ & FP8)

Deploying large models efficiently requires reducing precision from FP16 (16 bits) down to FP8 (8 bits) or INT4 (4 bits):

• Activation Outliers: Deep Transformer layers develop emergent outlier activation channels where a tiny fraction of channels ($0.1\%$) exhibit magnitudes up to $100\times$ higher than average. Clamping these destroys model perplexity!
• AWQ (Activation-Aware Weight Quantization): Protects salient weights corresponding to large activation channels, multiplying them by per-channel scales to balance dynamic range before INT4 quantization.
• GPTQ: Solves layer-wise second-order Taylor expansion using the inverse Hessian matrix $H^{-1}$: $$\Delta w = -\frac{w_i - \text{round}(w_i)}{[H^{-1}]_{ii}} H^{-1}_{:, i}$$
• Native FP8 (E4M3 & E5M2): Hardware-supported on Hopper Tensor Cores, delivering $2\times$ compute throughput and $2\times$ memory reduction with minimal calibration overhead.

Module 6.1: Speculative Decoding & Verification Kernels

The fundamental latency bottleneck of autoregressive generation is that producing $K$ tokens requires $K$ sequential forward passes through a gigantic model (e.g., 70B). However, validating $K$ candidate tokens simultaneously takes virtually the exact same execution time as generating 1 token on a modern GPU, because verification runs as a compute-bound parallel GEMM!

In Speculative Decoding:

1. A tiny, ultra-fast Draft Model (e.g. 7B or 1B) rapidly proposes a sequence of $\gamma$ speculative tokens.
2. The massive Target Model (e.g. 70B) executes a single parallel verification pass on all $\gamma$ tokens.
3. Tokens are accepted or rejected using modified rejection sampling: $$P(\text{accept } x) = \min\left(1, \frac{P_{\text{target}}(x)}{P_{\text{draft}}(x)}\right)$$

Because the sampling mathematics guarantees that the output probability distribution matches the target model exactly, speculative decoding achieves $2\times - 3\times$ lower latency with mathematical zero loss in output quality!

Module 6.2: Sparse Mixture-of-Experts (MoE) & Expert Parallelism

Scaling dense models beyond 100B parameters causes training and inference compute costs to skyrocket. Sparse Mixture-of-Experts (MoE) decouples parameter count from FLOPs per token by replacing the single dense Feed-Forward layer with $E$ independent expert networks (e.g., $E=8$ or $E=64$):

For each token, a gating network routes the vector to only the top-$k$ experts (e.g., $k=2$). A model can have 46.7 billion total parameters (like Mixtral 8x7B) while activating only 12.9 billion parameters per token, yielding the intelligence of a massive model with the inference speed of a small one!

To prevent the gating network from collapsing onto a single popular expert, architects enforce an auxiliary load-balancing loss $\mathcal{L}_{\text{aux}}$ across all experts:

$$\mathcal{L}_{\text{aux}} = \alpha \cdot E \sum_{i=1}^E f_i P_i$$

Module 6.3: Sub-Quadratic Sequence Models: Mamba & Selective State Spaces

Despite optimizations, Attention inherently retains $O(S^2)$ quadratic complexity with respect to context length. Selective State Space Models (SSMs), such as Mamba, replace attention with a continuous differential system discretized via Zero-Order Hold (ZOH):

By making the parameters $\Delta, B, C$ dynamic functions of the input token (selective mechanism) and utilizing a hardware-aware parallel associative scan in SRAM, Mamba achieves linear $O(S)$ computational complexity and zero KV cache memory growth during autoregressive decode!

Module 7.1: 3D Parallelism at Scale — TP, PP & ZeRO / FSDP

Training frontier foundation models (70B to 1T+ parameters) requires aggregating thousands of accelerator chips. Because no single GPU has sufficient memory or compute, architects combine three orthogonal dimensions of parallelism:

1. Tensor Parallelism (TP — Megatron-LM): Splits individual weight matrices within a layer across GPUs. In the MLP, the first matrix $W_1$ is partitioned by columns (requiring an All-Gather), and the second matrix $W_2$ is partitioned by rows (requiring an All-Reduce). Because TP introduces collective communication on every single layer, it requires ultra-high bandwidth NVLink ($900 - 1800 \text{ GB/s}$) and is strictly confined within a single 8-GPU node.
2. Pipeline Parallelism (PP): Partitions the model's $L$ layers across different server nodes. In the 1F1B (One-Forward-One-Backward) schedule, execution alternates between forward and backward micro-batches to bound activation memory while minimizing the idle pipeline bubble fraction: $$F_{\text{bubble}} = \frac{p - 1}{m + p - 1}$$
3. ZeRO / FSDP (Fully Sharded Data Parallelism): Shards the memory footprint of training states across data-parallel ranks:
   • ZeRO-1: Shards optimizer states ($4\times$ memory reduction).
   • ZeRO-2: Shards gradients ($8\times$ memory reduction).
   • ZeRO-3: Shards model parameters, streaming weights dynamically via All-Gather and releasing them immediately after computation.

Module 7.2: Interconnect Topologies & Scale-Out Networks

At cluster scale, network latency and bandwidth dictate training efficiency:

• Intra-Node (NVLink & NVSwitch): Provides $1.8 \text{ TB/s}$ bidirectional bandwidth per GPU, enabling zero-latency Megatron TP communication across the 8 GPUs on a baseboard.
• Inter-Node (InfiniBand NDR / RoCEv2): 400 Gbps to 800 Gbps per GPU. To avoid saturation on spine switches, engineers configure Rail-Optimized Leaf-Spine Topologies: GPU rank $i$ across all servers connects to the same leaf switch, allowing Data Parallel All-Reduce collectives to execute locally without crossing spine switches!
• Lossless Ethernet: Priority Flow Control (PFC) and Explicit Congestion Notification (ECN) prevent packet drops and buffer bloat that would otherwise cause catastrophic pipeline stalls.

Module 7.3: Chinchilla Scaling, MFU Optimization & Serving Economics

According to the Chinchilla Scaling Laws (Hoffmann et al.), compute-optimal training requires scaling model parameters $N$ and training tokens $D$ in equal proportion:

A 70B parameter model requires at least $1.4 \text{ Trillion tokens}$ to reach the compute-optimal frontier. However, production models (like LLaMA-3 with 15T tokens) are deliberately "overtrained" to minimize downstream per-token inference cost at scale!

Cluster efficiency is evaluated via Model FLOPs Utilization (MFU):

Achieving $>50\%$ MFU on a 16,384-GPU supercluster requires meticulous overlap of communication collectives with computation kernels, FP8 precision execution, and near-zero pipeline bubble overhead.