LLM Basics Series 4: Transformer Training and Inference

$p_\theta$ is a parametric conditional distribution. The subscript $\theta$ denotes every learnable parameter in the entire transformer: all $W_{Q}, W_{K}, W_{V}, W_{O}$ matrices across all heads and layers, all FFN weights, the embedding table, LayerNorm parameters, and the output head.

The output at each position is a categorical distribution over the vocabulary of size $V$. Here is how it is constructed:

1. The transformer processes tokens $x_{\leq t}$ through all $L$ blocks, producing $h_{t} \in \mathbb{R}^{d_{model}}$.
2. The output head projects to vocab size: $\text{logits}_{t} = h_{t} W_{\text{head}} \in \mathbb{R}^{V}$. This gives one raw score per vocabulary token.
3. Softmax converts logits to probabilities:

For each position, this gives a probability distribution over all $V$ tokens (e.g., 151,936 for Qwen3). The probabilities sum to 1.

This is a conditional probability: the distribution over the next token depends on all preceding tokens $x_{\leq t}$ through the autoregressive structure. Different contexts produce different distributions, all parameterized by the same $\theta$. Training adjusts $\theta$ via gradient descent on the cross-entropy loss so that $p_\theta$ assigns high probability to the actual next token in the training data.

Cross-entropy loss vs maximum likelihood. The pretraining objective above is written as maximizing log-likelihood. In practice, training frameworks minimize a loss. The two are the same objective with a sign flip:

• Maximum likelihood: $\max_\theta \sum_{t} \log p_\theta(x_{t} \mid x_{<t})$
• Cross-entropy loss: $\min_\theta -\sum_{t} \log p_\theta(x_{t} \mid x_{<t})$

The name "cross-entropy" comes from information theory. For a true distribution $q$ (one-hot: all probability mass on the actual next token $w^{\ast}$) and predicted distribution $p_\theta$, the cross-entropy is:

Since $q$ is one-hot, only one term survives:

This is exactly the negative log-likelihood of the true token. So "cross-entropy loss" and "negative log-likelihood" are the same formula when the target is a single token. Maximizing log-likelihood, minimizing negative log-likelihood, and minimizing cross-entropy all lead to the same gradient update.

End of expanded note.

Chunking long documents. A book or paper can be hundreds of thousands of tokens, far longer than the model's context window (e.g., 4K or 32K tokens). Long documents are split into chunks for training. Common approaches:

• Fixed-length chunking. Split into non-overlapping segments of the context window size (e.g., every 4,096 tokens). Simple but can cut in the middle of a sentence or paragraph, losing local coherence at chunk boundaries.
• Document-aware chunking. Split at natural boundaries (chapter breaks, section headers, paragraph breaks) and then pack segments up to the context length. This preserves local coherence. If a chapter is 10K tokens and the context window is 4K, split at the nearest paragraph break to 4K.
• Packing multiple short documents. Short documents (e.g., tweets, short articles) are concatenated into one training sequence with a special separator token (e.g., <|endoftext|>) between them, up to the context length. This avoids wasting compute on padding. The causal mask does not need modification because each document's tokens are simply earlier context for the next document's tokens; the model learns that the separator resets the topic.

Deduplication. Web crawls contain enormous amounts of duplicate content: the same Wikipedia article appears on dozens of mirror sites, the same news article is syndicated across hundreds of outlets, boilerplate text (cookie notices, navigation menus) appears on every page of a site. Training on duplicates wastes compute and biases the model toward memorizing repeated content. Deduplication operates at multiple levels:

• Exact deduplication. Hash each document (e.g., SHA-256 of the full text) and remove exact copies. Fast but misses near-duplicates (same article with a slightly different header).
• Near-deduplication with MinHash. Compute a set of hash signatures (MinHash) for each document based on its n-gram content. Two documents with a high fraction of matching signatures (e.g., Jaccard similarity above 0.8) are considered near-duplicates. One copy is kept, the rest are removed. This catches paraphrased or lightly edited copies. The LLaMA and Qwen technical reports describe using MinHash-based deduplication.
• Substring-level deduplication. Even after document-level deduplication, the same paragraph can appear across many different documents (e.g., a common definition that appears in multiple Wikipedia articles and textbooks). Suffix array-based methods (used in the "Deduplicating Training Data" paper by Lee et al.) find and remove repeated substrings across the entire corpus. This is more expensive but catches fine-grained repetition.

Quality filtering. Raw web crawls (Common Crawl) contain billions of pages, most of which are low quality: spam, SEO content, auto-generated text, navigation menus, error pages, and so on. Quality filtering keeps only pages that resemble "good" text (books, Wikipedia, well-written articles). Common methods:

• Heuristic rules. Remove documents that are too short, have too many special characters, have very low or very high word repetition ratios, contain too many URLs, or have a very low ratio of alphabetic characters. For example, a page where more than 30% of lines end with "..." is likely auto-generated spam.
• Language detection. Use a language classifier (e.g., fastText) to keep only documents in the target language(s). This removes garbled text, mixed-language spam, and encoding errors.
• Perplexity filtering. Train a small language model (e.g., a KenLM n-gram model) on a known high-quality corpus (Wikipedia, books). Score each web page by its perplexity under this model. Pages with very high perplexity (the model finds them surprising and incoherent) are likely low quality and are removed. Pages with very low perplexity might be too repetitive or templated. A middle range is kept. This is the approach used in CCNet (the Common Crawl filtering pipeline used by LLaMA).
• Classifier-based filtering. Train a binary classifier (e.g., a small BERT or logistic regression on TF-IDF features) to distinguish "Wikipedia-like" text from random web text. Score each page and keep those above a threshold. GPT-3 used a classifier trained on WebText (Reddit-upvoted links) as positive examples and raw Common Crawl as negative examples.

Toxic and harmful content removal. Even after quality filtering, the corpus can contain hate speech, explicit content, personally identifiable information (PII), and other harmful material. Approaches include:

• Keyword and regex filters. Remove documents containing known slurs, explicit terms, or patterns matching PII (phone numbers, email addresses, social security numbers). Fast but high false-positive rate (the word "kill" appears in legitimate contexts) and misses rephrased toxicity.
• Toxicity classifiers. Use a trained classifier (e.g., Perspective API, or a custom model trained on labeled toxic/non-toxic data) to score each document. Remove or downweight documents above a toxicity threshold. More robust than keyword matching but not perfect.
• PII scrubbing. Use named entity recognition and regex patterns to detect and replace names, addresses, phone numbers, and other PII with placeholder tokens. This is especially important for compliance with privacy regulations.
• Domain blocklists. Maintain lists of known harmful, adult, or spam domains and exclude all pages from those domains.

Data mixing and weighting. The final training corpus is a mixture of sources with intentional proportions. A typical mix might be:

• Web text (filtered Common Crawl): ~60-70%
• Code (GitHub): ~5-10%
• Books: ~5-10%
• Wikipedia: ~3-5%
• Scientific papers (ArXiv): ~3-5%
• Curated/instruction data: ~1-5%

Higher-quality sources (books, Wikipedia) are often upsampled (repeated multiple times) despite being smaller, because their quality-per-token is much higher than raw web text. The mixing ratios are tuned empirically and can significantly affect downstream performance.

End of expanded note.

Before any training or inference, raw text must be converted into integer token IDs that the embedding table can look up. This is done by a tokenizer, which is trained separately from the model.

Byte-Pair Encoding (BPE) is the most common tokenization algorithm (used by GPT, LLaMA, Qwen). It works by:

1. Start with individual characters (or bytes) as the initial vocabulary.
2. Count all adjacent pairs in the training corpus.
3. Merge the most frequent pair into a new token.
4. Repeat until the vocabulary reaches the target size (e.g., 32K, 50K, 152K tokens).

For example, if "th" appears very frequently, it becomes a single token. Then "the" might be merged next. Common words like "the" become single tokens, while rare words are split into subword pieces: "unforgettable" might become ["un", "forget", "table"] or ["un", "forg", "ettable"] depending on the learned merges.

Why tokenization matters:

• Different models use different tokenizers. The same text produces different token counts and different token IDs across models. Qwen3 with 152K vocab tokenizes text more efficiently (fewer tokens per sentence) than GPT-2 with 50K vocab.
• The vocabulary size $V$ determines the embedding table size ($V \times d_{model}$ parameters) and the output head size.
• Tokenization affects context window utilization: more efficient tokenization fits more text into the same context length.

End of expanded note.

Modern LLM training uses mixed precision to reduce memory and increase throughput. The key idea: use lower precision (16-bit) for most computations and higher precision (32-bit) only where it matters.

Three formats:

• FP32 (32-bit float): full precision. 1 sign bit, 8 exponent bits, 23 mantissa bits. Used for optimizer states ($m_{t}$, $v_{t}$ in AdamW) and the master copy of weights.
• FP16 (16-bit float): half precision. 1 sign bit, 5 exponent bits, 10 mantissa bits. Smaller range ($\pm 65504$), which can cause overflow during training.
• BF16 (bfloat16): 1 sign bit, 8 exponent bits, 7 mantissa bits. Same range as FP32 (8 exponent bits) but lower precision (7 vs 23 mantissa bits). Preferred for LLM training because it avoids overflow issues.

The mixed precision recipe:

1. Store a master copy of weights in FP32.
2. Cast weights to BF16 for the forward and backward pass (half the memory, faster matmuls on modern GPUs).
3. Compute gradients in BF16.
4. Update the FP32 master weights using the FP32 optimizer states.

This gives nearly the same training dynamics as full FP32, at roughly half the memory for model weights and significantly faster computation. The FP32 optimizer states (AdamW stores $m_{t}$ and $v_{t}$ per parameter) are the dominant memory cost.

End of expanded note.

Default choice: BF16. If your hardware supports it (NVIDIA A100, H100, or newer; Google TPUs), use BF16 for training. BF16 has the same exponent range as FP32 (8 exponent bits), which means it can represent very large and very small values without overflow or underflow. This makes it stable for LLM training without requiring loss scaling. Nearly all large-scale pretraining runs since 2022 use BF16.

When to use FP16. Older GPUs (V100, T4, consumer GPUs before RTX 30-series) do not have native BF16 support. On these GPUs, FP16 is the only half-precision option. FP16 has a narrower range ($\pm 65504$) due to only 5 exponent bits, which means gradients or activations can overflow during training. To compensate, FP16 training requires loss scaling: multiply the loss by a large factor (e.g., 1024) before the backward pass to push small gradients into FP16's representable range, then divide the gradients by the same factor before the optimizer step. PyTorch's GradScaler handles this automatically. FP16 with loss scaling works well for fine-tuning and smaller models, but can be fragile for large-scale pretraining where activation magnitudes vary widely across layers.

When to use full FP32. Two situations: (1) optimizer states, which should always be in FP32 regardless of the forward/backward precision, because the small weight updates from AdamW can vanish in half-precision arithmetic; (2) debugging training instability, where switching to full FP32 helps isolate whether a NaN or divergence is caused by precision issues or by a genuine optimization problem (bad learning rate, data corruption, etc.).

Inference precision. For inference (no gradient computation), precision requirements are more relaxed. BF16 is the standard serving format. FP16 also works well for inference since there is no gradient accumulation to overflow. For further memory savings, quantized formats (INT8, INT4) are increasingly common for inference and can cut memory by another 2-4x with modest quality loss. The choice depends on whether latency, throughput, or quality is the priority.

End of expanded note.

A pretrained model is trained to predict the next token in web text. If you give it the prompt "What is the capital of France?", it might continue with "What is the capital of Germany? What is the capital of Italy?" because in its training data, quiz questions often appear in lists. It is completing the document, not answering the question.

SFT teaches the model that when it sees a question, the appropriate continuation is an answer, not more questions. RL alignment then refines the style, safety, and helpfulness of those answers based on human preferences.

This three-stage pipeline (pretrain → SFT → alignment) is the standard recipe used by GPT-4, Claude, LLaMA-Chat, Qwen-Chat, and most production LLMs.

End of expanded note.

When you need a model that handles multiple domains (e.g., medical and legal), there are several approaches to combine fine-tuned capabilities without losing either.

Approach 1: multi-task fine-tuning. Mix all domain data together in one SFT stage, then one RL stage. You control the domain balance through data mixing ratios (e.g., 50% medical, 50% legal). The model learns both domains jointly and there is no forgetting. This is the simplest and most common approach when you have access to all data at once.

Approach 2: sequential fine-tuning. Fine-tune one domain after another: base → SFT medical → SFT legal → RL. The problem is catastrophic forgetting: when you fine-tune on legal data, gradient updates overwrite the medical-specialized weights, and medical performance degrades. You can mitigate this by mixing in some medical data during the legal stage (experience replay), but the balance is fragile to tune.

Approach 3: model merging (weight averaging). Fine-tune separate models from the same base, then merge their weights arithmetically:

This works surprisingly well when both models started from the same base, because the fine-tuned weights remain in a similar region of parameter space. More advanced merging methods include TIES-Merging (only merge parameters that changed significantly, resolve sign conflicts), DARE (randomly drop small weight deltas before merging to reduce interference), and SLERP (spherical interpolation on the parameter hypersphere). Model merging is cheap (no training, just arithmetic) and popular in the open-source community. The limitation is that it is heuristic: there is no guarantee the merged model is optimal.

Approach 4: LoRA adapters. Train lightweight low-rank adapters per domain on a shared frozen base. Each adapter modifies only ~0.1% of parameters. At inference, load the appropriate adapter for the detected domain, or merge multiple adapters into the base weights: $W = W_{\text{base}} + \Delta W_{\text{medical}} + \Delta W_{\text{legal}}$. This avoids catastrophic forgetting because the base is frozen. Each adapter is small and fast to train. The downside is that independently trained adapters may interfere when combined.

Approach 5: Mixture of Experts (MoE). MoE is the most principled architectural approach. Rather than merging after the fact, the model is designed from the start with multiple specialist sub-networks (experts) and a learned router that selects which experts to activate for each input.

How MoE works inside a transformer block. In a standard transformer, the FFN is a single two-layer MLP applied to every token. In an MoE transformer, the FFN is replaced by $E$ parallel expert networks (each is a separate FFN with its own weights) plus a gating router:

where $g_{i}(x)$ is the gating weight for expert $i$, computed by the router:

The router is a simple linear layer $W_{\text{gate}} \in \mathbb{R}^{d_{model} \times E}$ that produces a score for each expert. TopK selects only the top $k$ experts (e.g., $k = 2$ out of $E = 64$) and zeros out the rest. Only the selected experts run their FFN computation, so the actual compute per token is much smaller than having all $E$ experts active.

Concrete example: Qwen3-235B-A22B. This model has 235B total parameters but only 22B activated per token. Each transformer block has 128 expert FFNs, and the router selects 8 of them for each token. Different tokens in the same sequence can activate different experts. A medical question might route to experts that specialize in scientific reasoning, while a legal question in the same batch might route to different experts that specialize in formal language and citation patterns.

Why MoE helps with multi-domain capability. In a dense model, every parameter is used for every token. Specializing in medical language and legal language must share the same FFN weights, creating tension. In an MoE model, different experts can specialize in different domains or capabilities. The router learns to direct medical tokens to medical-specialized experts and legal tokens to legal-specialized experts, without interference. This specialization emerges naturally during training without explicit domain labels.

The load balancing problem. Without constraints, the router might learn to send all tokens to the same few experts (winner-take-all), leaving most experts unused. This wastes parameters and defeats the purpose. MoE training adds an auxiliary loss that penalizes imbalanced routing:

where $f_{i}$ is the fraction of tokens routed to expert $i$ and $p_{i}$ is the average router probability assigned to expert $i$. This loss is minimized when all experts receive equal traffic. It is weighted by a small coefficient (e.g., 0.01) and added to the main next-token prediction loss.

Trade-offs of MoE.

• Parameter efficiency. MoE can scale total parameters (and therefore knowledge capacity) without proportionally scaling compute per token. Qwen3-235B has 10x the parameters of Qwen3-32B but only ~4x the active parameters per token.
• Memory cost. All expert weights must be stored in memory even though only a few are active per token. Qwen3-235B needs enough GPU memory to hold 235B parameters, not just 22B.
• Communication overhead. In distributed training, different experts may live on different GPUs. The router must send tokens to the right GPU, which adds communication cost (all-to-all communication).
• Training stability. Router training can be unstable early on. The load balancing loss, careful initialization of the gate, and sometimes expert dropout are needed to stabilize training.

Is MoE the dominant approach? MoE is increasingly adopted for large-scale models (Mixtral, Qwen3-235B, GPT-4 is widely believed to be MoE, DeepSeek-V2/V3). It is the most principled way to scale knowledge capacity while keeping inference cost manageable. However, for smaller models or when adding a new domain to an existing dense model, multi-task mixing or LoRA adapters are more practical because they do not require changing the architecture. MoE is best thought of as an architectural choice made before pretraining, not a post-hoc merging technique.

Summary.

For most cases: multi-task mixing if you can combine the data upfront, model merging or LoRA if you need to add domains to an existing model without retraining, and MoE if you are designing a large-scale system from scratch where multi-domain capacity and inference efficiency both matter.

End of expanded note.

Consider three tokens with logits $[3.0, 1.0, 0.5]$.

At $\tau = 1.0$: softmax gives approximately $[0.78, 0.11, 0.06]$. Token 1 dominates.

At $\tau = 0.5$: logits become $[6.0, 2.0, 1.0]$. Softmax gives approximately $[0.97, 0.02, 0.01]$. Token 1 almost certain.

At $\tau = 2.0$: logits become $[1.5, 0.5, 0.25]$. Softmax gives approximately $[0.49, 0.18, 0.14]$. Much more uniform.

Dividing by $\tau$ scales the logit differences. Small $\tau$ amplifies differences (sharp distribution), large $\tau$ compresses differences (flat distribution). The ranking of tokens does not change; only the confidence does.

End of expanded note.

Typical ranges. Common values of $k$ fall between 10 and 100. Smaller $k$ (10-20) keeps generation tightly focused on the most probable tokens, producing coherent but less diverse text. Larger $k$ (50-100) allows more variety but increases the chance of sampling an off-topic or awkward token. Values above 100 rarely help because the probability mass in the tail is negligible for most positions.

Interaction with temperature. Temperature and top-k interact in important ways. High temperature flattens the distribution, pushing probability mass toward lower-ranked tokens. If you use high temperature ($\tau = 1.0$+) with a large $k$, the model can sample from tokens that were originally very unlikely, leading to incoherent output. The safe combination is: if you raise temperature, lower $k$ (or use top-p instead). If you lower temperature, a larger $k$ is harmless because the distribution is already sharp and the extra candidates receive negligible probability anyway.

When to prefer top-p over top-k. The fundamental limitation of top-k is that it ignores the shape of the distribution. At a position where the model is confident (one token has 0.9 probability), $k = 50$ includes 49 tokens that collectively share 0.1 probability, adding noise. At a position where 20 tokens are roughly equally likely, $k = 10$ cuts off good candidates. Top-p ($p = 0.9$ or $0.95$) adapts to both cases automatically. For this reason, most production LLM APIs default to top-p rather than top-k. If you do use top-k, $k = 40$-$50$ with $\tau = 0.7$-$0.9$ is a reasonable starting point for general text generation.

End of expanded note.

Top-k's problem: $k$ is fixed regardless of context. At some positions the model is very confident (one token has 0.9 probability), but top-k still samples from $k$ tokens, introducing unnecessary noise. At other positions many tokens are plausible, but $k$ might cut off good candidates.

Top-p adapts: it includes however many tokens are needed to cover probability mass $p$. Confident positions use fewer tokens, uncertain positions use more.

In practice, top-p ($p = 0.9$ or $0.95$) is the more common default in production LLMs. Many systems combine temperature + top-p: first adjust sharpness with temperature, then truncate with top-p. Some also combine top-k and top-p together (apply both filters).

Typical settings for different use cases:

• Code generation: low temperature (0.2–0.4), low top-p (0.9). Correctness matters more than creativity.
• Creative writing: higher temperature (0.7–1.0), higher top-p (0.95). Diversity and surprise are desirable.
• Factual Q&A: temperature close to 0 or greedy. Minimize hallucination.

End of expanded note.

Training context length. The context window used during pretraining determines what the model can reliably handle at inference. Training at 4K means the model has never seen attention patterns spanning more than 4K positions, so inference at 32K will degrade without additional work. The cost tradeoff is direct: doubling the training context length roughly quadruples the attention compute (the $T^2$ factor) and doubles the KV cache memory. Most teams start pretraining at a shorter context (2K-4K) and extend later, because the majority of training data consists of short documents and the model learns most of its language ability from local patterns.

Progressive context extension. Modern LLMs typically use a staged approach: pretrain at 4K-8K for most of training, then do a short continued-pretraining phase (1-5% of total tokens) at the target context length (32K, 128K). This is far cheaper than training at 128K from the start. The continued-pretraining phase uses long documents (books, concatenated articles, code repositories) so the model actually sees and learns to use the extended positions. LLaMA 3.1 and Qwen3 both adopted this approach. The RoPE base frequency is typically increased during extension (see the RoPE frequency base note in Series 3).

Inference considerations. A model trained at 32K can be served at 32K, but users may not always need the full window. Shorter inputs are cheaper (less KV cache memory, faster attention). When serving many concurrent users, the effective batch size is often limited by KV cache memory, not compute. Reducing the maximum generation length or using a shorter system prompt can significantly increase throughput. For applications that need to process very long documents (100K+ tokens), consider whether the task genuinely requires full attention over the entire document or whether a chunked approach (process segments independently, then combine) would suffice at much lower cost.

Practical defaults. For fine-tuning an existing model, match or stay below the base model's training context length unless you also do RoPE scaling and continued pretraining on long-context data. For new pretraining, 4K-8K for the main phase and 32K-128K for the extension phase is the current standard. Choose the final context length based on the target application: 4K is sufficient for single-turn Q&A, 32K covers most document analysis tasks, and 128K+ is needed for book-length inputs or multi-document reasoning.

End of expanded note.

Step-by-step. Generating the sequence "The cat sat":

Step 1: process prompt "The".

• Compute $q_{0}, k_{0}, v_{0}$ for "The".
• Cache: $K = [k_{0}]$, $V = [v_{0}]$.
• Attention: $q_{0}$ attends to $[k_{0}]$. Output predicts "cat".

Step 2: generate "cat".

• Compute $q_{1}, k_{1}, v_{1}$ for "cat" only (one token, not the full sequence).
• Append to cache: $K = [k_{0}, k_{1}]$, $V = [v_{0}, v_{1}]$.
• Attention: $q_{1}$ attends to $[k_{0}, k_{1}]$. Output predicts "sat".

Step 3: generate "sat".

• Compute $q_{2}, k_{2}, v_{2}$ for "sat" only.
• Append to cache: $K = [k_{0}, k_{1}, k_{2}]$, $V = [v_{0}, v_{1}, v_{2}]$.
• Attention: $q_{2}$ attends to $[k_{0}, k_{1}, k_{2}]$. Output predicts the next token.

At each step, only one new token is processed through the Q, K, V projections and FFN. The expensive part (projecting and transforming all previous tokens) is not repeated.

Memory cost. The KV cache stores $K$ and $V$ for every layer and every head. For a model with $L$ layers, $H_{kv}$ KV heads, per-head dimension $d_{k}$, and sequence length $T$:

For Qwen3-8B ($L = 36$, $H_{kv} = 8$ with GQA, $d_{k} = 128$) at $T = 32{,}768$ in BF16 (2 bytes):

This is per sequence. Serving many concurrent users requires multiplying by the number of active sequences, which is why KV cache memory is often the bottleneck for inference serving.

GQA reduces KV cache. Grouped-query attention (discussed in Series 3) shares K and V across groups of query heads. Qwen3-8B has 32 query heads but only 8 KV heads, reducing the cache by 4x compared to full multi-head attention.

End of expanded note.

The scaling law. The cross-entropy loss $L$ on held-out text follows approximately:

where $N$ is the number of parameters, $D$ is the number of training tokens, and $L_\infty$ is an irreducible loss (entropy of natural language). The exponents $\alpha_{N} \approx 0.076$ and $\alpha_{D} \approx 0.095$ mean that both axes give diminishing returns: doubling model size or data does not halve the loss.

Compute-optimal training (Chinchilla). For a fixed compute budget $C \propto N \times D$ (FLOPs $\approx 6ND$ for a forward + backward pass), the optimal allocation is roughly:

Model size and data should scale at the same rate. The original GPT-3 (175B parameters, 300B tokens) was undertrained by this criterion. Chinchilla (70B parameters, 1.4T tokens) matched GPT-3 performance with less compute by using a better ratio.

Modern LLMs (LLaMA 3, Qwen3) train well beyond the Chinchilla-optimal point, because inference cost depends on model size but not training data. A smaller model trained on more data is cheaper to deploy even if training costs more.

End of expanded note.

The problem with standard attention. Standard attention computes the full $T \times T$ score matrix, writes it to GPU main memory (HBM), applies softmax, writes the result back, then multiplies by $V$. For $T = 32{,}768$, this matrix has about 1 billion entries (~2 GB in BF16 per head per layer). The computation is not limited by arithmetic (GPUs have plenty of FLOPS); it is limited by memory bandwidth, the time spent reading and writing this huge intermediate matrix.

GPU memory hierarchy. A GPU has two levels of memory:

• SRAM (on-chip): ~20 MB, ~19 TB/s bandwidth. Very fast, very small.
• HBM (off-chip): ~80 GB, ~2 TB/s bandwidth. Large, but ~10x slower.

Standard attention writes the full $T \times T$ matrix to HBM and reads it back multiple times. The GPU spends most of its time waiting for data transfers, not doing math.

Flash Attention's solution: tiling. Instead of computing the full $T \times T$ matrix at once, Flash Attention processes Q, K, V in small blocks (tiles) that fit entirely in SRAM. It never writes the full attention matrix to HBM:

Standard:  Q, K → [full T×T matrix in HBM] → softmax → [full T×T in HBM] → ×V → output

Flash:     for each block of Q rows:
             for each block of K, V columns:
               compute scores (in SRAM)
               update running softmax (in SRAM)
               accumulate output (in SRAM)
             write final output block to HBM

The hard part: incremental softmax. Softmax needs the entire row to normalize: $\text{softmax}(s_{j}) = e^{s_{j}} / \sum_{j'} e^{s_{j'}}$. If you process K in blocks, you do not have all $T$ scores at once. Flash Attention uses the online softmax trick: maintain a running maximum $m$ and running sum $\ell$ as each block is processed, and rescale partial results when the maximum changes.

For each new block $b$:

1. Compute local scores: $s^{(b)} = q \cdot K_{b}^\top / \sqrt{d_{k}}$.
2. Update running max: $m_{\text{new}} = \max(m_{\text{old}}, \max(s^{(b)}))$.
3. Rescale old accumulator: $O \leftarrow O \cdot e^{m_{\text{old}} - m_{\text{new}}}$.
4. Accumulate: $O \leftarrow O + e^{s^{(b)} - m_{\text{new}}} \cdot V_{b}$.
5. Update running sum: $\ell \leftarrow \ell \cdot e^{m_{\text{old}} - m_{\text{new}}} + \sum e^{s^{(b)} - m_{\text{new}}}$.

After all blocks: $O \leftarrow O / \ell$. The rescaling in step 3 is the key: when a new block has larger scores than all previous blocks, the old partial results are scaled down to match. The final result is mathematically identical to computing softmax over the full row.

What you gain:

Flash Attention 2 and 3. Flash Attention 2 improved parallelism by distributing work across GPU thread blocks more efficiently (parallelize over the sequence length dimension, not just batch and head). Flash Attention 3 further optimizes for newer GPU architectures (Hopper). All modern LLM training and inference frameworks (PyTorch 2.0+, vLLM, TensorRT-LLM) use Flash Attention by default.

End of expanded note.

LLMs have no persistent memory across conversations. When you start a new conversation, the model knows nothing about previous conversations. Its only information comes from (1) what it learned during pretraining and fine-tuning (baked into the weights $\theta$), and (2) what you provide in the current context window.

The system prompt. Most chat APIs allow a system prompt (or system message) that is prepended to every conversation. This is the primary way to give the model persistent instructions or knowledge within a session:

System: You are a medical assistant specializing in cardiology. 
        Always cite clinical guidelines when answering.
        The patient's history: [... relevant details ...]

User:   What are the treatment options for atrial fibrillation?

The system prompt is just regular tokens in the context window. It has no special mechanism; it works because the model attends to it on every turn. If the system prompt is 2,000 tokens and the context window is 32K, you have 30K tokens left for conversation.

Retrieval-Augmented Generation (RAG). For knowledge that does not fit in the system prompt or changes frequently, RAG retrieves relevant documents from an external database and inserts them into the prompt before the model generates:

1. User asks a question.
2. A retrieval system (embedding similarity search, keyword search, etc.) finds the most relevant documents from a knowledge base.
3. The retrieved documents are inserted into the prompt as context.
4. The model generates an answer grounded in the provided documents.

This is how most production systems handle large or dynamic knowledge bases (company docs, product catalogs, legal databases). The model does not "remember" the documents; they are provided fresh in each query's context window.

Fine-tuning. To permanently bake domain knowledge into the model's weights, fine-tune on domain-specific data (continued pretraining or SFT). This is more expensive than RAG but makes the knowledge available without consuming context window space. The trade-off: fine-tuning is slow to update (retrain for new information), while RAG can be updated instantly by changing the document database.

Practical guidelines:

• For behavioral instructions (tone, format, persona): use the system prompt.
• For specific facts or documents the model needs per query: use RAG.
• For broad domain competence (the model should "think like a doctor"): fine-tune.
• For conversation history: include previous turns in the context. When the conversation exceeds the context window, older turns must be truncated or summarized.

End of expanded note.

The prompt is the only input the model sees. Small changes in phrasing can cause large changes in output quality. Some practical guidelines:

General principles:

• Be specific. "Summarize this article" is vague. "Summarize this article in 3 bullet points, each under 20 words, focusing on the methodology" tells the model exactly what you want.
• Provide examples (few-shot). Showing 2–3 examples of the desired input-output format is often more effective than describing the format in words. The model pattern-matches from examples.
• Specify the output format. If you want JSON, say "respond in valid JSON with keys: name, age, diagnosis." If you want a table, show the header row. Ambiguous format instructions produce ambiguous output.
• Assign a role. "You are an experienced cardiologist reviewing a case" activates different knowledge patterns than "Answer this medical question." Roles prime the model toward domain-appropriate language and reasoning depth.

Task-specific guidelines:

• Factual Q&A. Ask the model to cite sources or say "I don't know" if uncertain. Include "Answer based only on the following context: [...]" to reduce hallucination. Use low temperature.
• Code generation. Specify the language, framework, and version. Provide function signatures or test cases as constraints. Ask the model to think step-by-step before writing code.
• Summarization. Specify length, audience, and focus. "Summarize for a technical audience in 100 words" vs "Explain to a 10-year-old in 2 sentences" produce very different outputs from the same input.
• Analysis and reasoning. Use chain-of-thought prompting: "Think step by step before giving your final answer." This significantly improves accuracy on math, logic, and multi-step reasoning tasks.
• Creative writing. Provide constraints (genre, tone, length, audience) rather than leaving it open-ended. "Write a noir detective monologue, 200 words, first person" gives better results than "Write something creative."
• Classification and extraction. Define the categories explicitly and provide one example per category. For extraction, specify the exact fields you want and their types.

Common prompt mistakes:

• Too vague: "Make this better." Better how? More concise? More formal? More accurate?
• Contradictory instructions: "Be concise and thorough." Pick one priority.
• Assuming prior context: "As I said earlier..." in a new conversation. The model has no memory of earlier conversations.
• Over-constraining: 20 instructions in the system prompt can cause the model to ignore some. Prioritize the most important 3–5 constraints.

End of expanded note.