C. Answers to check questions and stretch exercises

Try each question or exercise before reading its answer. Each chapter below keeps the original question and exercise numbering. Check answers explain the concepts; stretch answers give worked results or a solution approach and verification criteria. Measurements and plots depend on your checkpoint, data, seed and hardware: the guidance below describes what to measure and what to expect, rather than claiming unrun experimental results. If an answer surprises you, reread the corresponding chapter section.

1. The big picture

Check your understanding

Each token is chosen from logits that depend on all previous tokens, including the ones just generated. Token 2 can’t be computed until token 1 exists, so 40 tokens need 40 dependent model calls (speculative decoding, Chapter 26, bends this rule by guessing).
Both read the weights once, but prefill does 1,000 tokens’ worth of arithmetic per weight read, so it’s compute-bound and uses the tensor cores fully. One decode step does one token’s arithmetic per weight read and is memory-bound. The 1,000× extra arithmetic costs far less than 1,000× the time.
7B × 2 bytes = 14 GB per token; 1 TB/s ÷ 14 GB ≈ 71 tokens/s at most. At 4 bits (3.5 GB plus scales) the ceiling rises to about 270-285 tokens/s.

Stretch exercises

Sampling at temperature 1.2 gives varied continuations and can choose less likely tokens; greedy decoding repeats the same choices for fixed logits. The prompt, weights, vocabulary and autoregressive dependency stay the same. Three sampled runs can still coincide, especially with a fixed seed.
With warmed-up, synchronized runs and EOS disabled for this measurement, let the timings be $t_1$ and $t_{101}$. TTFT is approximately $t_1$, and mean decode time is $(t_{101}-t_1)/100$. Prefill usually dominates short answers; enough decode steps dominate long ones.
FP32 doubles weight bytes relative to BF16, so an otherwise identical bandwidth-bound decode should take about twice as long. Actual slowdown depends on kernels, bandwidth and launch overhead; prefill also depends on the hardware’s FP32 compute path.

2. Tensors

Check your understanding

Matmul contracts the last axis of the left operand (4) with the first axis of the right (4) and broadcasts the leading batch axis (2); the remaining axes 3 and 5 form the output.
No: contiguous strides for [3, 2] are (2, 1). Strides (1, 3) are what you get from transposing a contiguous [2, 3] tensor (.T), which swaps shape and strides without moving data.
A square matrix and its transpose have the same shape, so every shape check passes while the computation is wrong. Only a value check against a reference catches it.
28 × 2 × 8 × 4,096 × 128 × 2 bytes = 469,762,048 bytes = 448 MiB.

Stretch exercises

The shapes are [8, 7, 6, 5], [3, 3] and [4, 4, 5], respectively. Align broadcast dimensions from the right; the final expression forms every pairwise difference between the four rows.
Transpose swaps entries a and b in both shape and strides, leaving the offset unchanged. For a valid nonnegative start and positive step, slicing sets the offset to offset + start * strides[axis], the stride to step * strides[axis], and the length to max(0, ceil((shape[axis] - start)/step)). Normalize slice bounds first; PyTorch does not support negative-step tensor slices.
Measure the ratio of median loop time to median @ time with the same dtype and thread count. Python dispatch and scalar tensor operations dominate the loop, while @ runs optimized compiled code; there is no machine-independent slowdown factor.
Broadcast the leading shapes, allocate [..., M, N], and loop over each output batch index and (m, n, k). Map an output batch coordinate to 0 whenever an operand’s corresponding batch dimension is 1, including implicit leading dimensions. Compare values with @ on unequal batch ranks and singleton dimensions.

3. How networks learn

Check your understanding

The gradient points in the direction of steepest increase of the loss; stepping against it decreases the loss fastest for a small step.
A value used in several places receives gradient contributions from each use (the multivariable chain rule sums them). Assigning with = would keep only the last contribution.
Without nonlinearities, a stack of linear layers collapses into one linear map, so depth adds nothing: the network can only represent linear functions.
eval() changes module behavior (dropout off, normalization uses running statistics). It doesn’t stop autograd from recording the graph; no_grad() (or inference_mode()) does, saving memory and time.
About 16-18 bytes per parameter (weights, master copy, gradients, two AdamW moments) against 2 for BF16 inference: roughly 8-9 times as much, before activations.

Stretch exercises

Compute s = sigmoid(x.data), create a result with value s, and have its backward closure add s * (1-s) * out.grad to x.grad. Use a stable sigmoid evaluation and compare with central differences at negative, zero and positive inputs.
With identical initialization, 0.005 usually progresses slowly, 0.05 often converges sooner, and 0.5 can oscillate or fail. These are expectations, not guaranteed outcomes: seed, loss and saturation affect the curves. Plot all runs over the same number of updates.
Momentum carries a running gradient, for example v = mu*v + grad; p -= lr*v. Adam tracks first and second moments, bias-corrects both, then updates with lr*m_hat/(sqrt(v_hat)+eps). Clear accumulated gradients each step and tune each optimizer fairly; steps to convergence depend on initialization and hyperparameters.
Use a nonlinear hidden layer between two nn.Linear modules and train on the four XOR pairs with AdamW. For 1B FP32 parameters, weights, gradients and two FP32 moments total about 16 GB, excluding activations and small step counters. A mixed-precision layout with BF16 weights/gradients and FP32 master weights/moments also totals 16 GB; FP32 gradients raise it to 18 GB. BF16 inference weights need 2 GB.

4. Text as numbers

Check your understanding

Byte-level BPE starts from all 256 byte values, so any string, in any language or with any emoji, can be written as bytes and therefore as tokens. A word tokenizer has no ID for words it never saw.
Later merges were learned on text where earlier merges had already been applied; they refer to tokens that only exist after those merges. Applying them in another order produces different (and unknown-to-the-model) token sequences.
The model was trained to begin answers after <|im_start|>assistant\n. Without it, the most likely continuation is often more of the user’s turn or a new role marker, so the model continues the prompt instead of answering.
Overlapping windows share text. If one window is in training and its neighbor in validation, the model has already seen most of the validation text, and validation loss underestimates the true error.
An embedding maps each token ID to a vector regardless of where it appears. A permuted sequence gets the same set of vectors; order must come from positional embeddings or RoPE.

Stretch exercises

Counts depend on the learned merge table: case, spaces and punctuation change which pairs were frequent enough to merge. Encode all three with the same trained tokenizer and print IDs and decoded pieces; no fixed counts follow from vocabulary size alone.
Divide the story’s UTF-8 byte count by its token count for each trained vocabulary. More merges normally increase bytes per token on the training text, with diminishing gains once common substrings are represented. The knee depends on the corpus; evaluate held-out text too.
Print both tokenizers’ decoded pieces for exactly the same paragraph. GPT-2’s larger vocabulary often preserves common words or leading-space word pieces that the 512-token tokenizer splits, but its pre-tokenizer also differs, so vocabulary size alone does not explain every difference.
Keep live symbols in a linked list and heap entries (merge_rank, position, version). Pop the best valid pair, merge it, invalidate stale entries and enqueue only its new neighbours. Preserve rank and left-to-right tie ordering, verify identical IDs, then benchmark the whole story.

5. Attention

Check your understanding

Each query asks “how much should I read from each key?”, and its weights over the keys must form a distribution (sum to 1). That’s a normalization along each row, over keys.
Masked positions must get exactly zero weight. Setting their scores to $-\infty$ before the softmax does that and renormalizes the rest; zeroing weights after the softmax would leave the row not summing to 1.
Only the token at position 0 (its query, key and value). Causal masking hides every later position.
Dot products of $d$-dimensional random vectors have variance proportional to $d$. Without the scaling, scores are large, the softmax saturates to nearly one-hot, and gradients through it vanish.
The group size is 16 / 8 = 2, so query head 11 reads KV head 11 // 2 = 5.

Stretch exercises

Compute scores = x @ x.T, mask entries above the diagonal to negative infinity, and apply row softmax with scale 1. The causal weight table, rounded to four decimals, is:

1.0000  0       0       0       0       0
0.3680  0.6320  0       0       0       0
0.2284  0.3893  0.3822  0       0       0
0.2046  0.2956  0.2915  0.2084  0       0
0.1753  0.2250  0.2269  0.1570  0.2158  0
0.1385  0.2184  0.2128  0.1420  0.0988  0.1896

Before rounding, every row sums to 1 and has zeros at future positions.

Fill [B, T, H*d] with distinct values encoding position, head and feature. Splitting must reshape to [B, T, H, d] and transpose to [B, H, T, d]; assert individual head contents as well as the correct round trip. A wrong split and its matching wrong inverse can pass a round-trip check together.
For absolute query/key positions, use allowed = (key_pos <= query_pos) & (key_pos > query_pos - w). This includes the current key and at most w-1 earlier keys; handle queries with no valid keys explicitly.
The score/probability tensors each have B*H*T*T elements, so doubling T roughly quadruples their memory and attention arithmetic. Reset peak allocation statistics and synchronize each measurement; estimate the fit limit from measured peak memory, leaving room for weights and other tensors. Small-size timing may be dominated by overhead.

6. The transformer block

Check your understanding

The MLP mixes features within each position; attention mixes information across positions.
A pre-norm block computes x + attn(norm(x)) and x + mlp(norm(x)): if both branch outputs are zero, each addition returns x unchanged. That’s why residual networks start close to the identity and train stably.
Storage: one vocabulary-by-width matrix instead of two (for GPT-2, 38.6M fewer parameters). Training: the matrix gets gradients from both uses, so input and output representations are coupled.
LayerNorm normalizes and then applies a learned scale $\gamma$ and shift $\beta$, which can move the output to any mean and variance.
Attention has four $D \times D$ projections (Q, K, V, output): $4D^2$. The MLP has $D \times 4D$ and $4D \times D$: $8D^2$. Total $12D^2$, ignoring biases and norms.

Stretch exercises

With vocabulary 50,257, context 1,024 and a tied head, the count is $50{,}257(1{,}024)+1{,}024^2+24(12(1{,}024)^2+13(1{,}024))+2(1{,}024)=354{,}823{,}168$. Thus “355M” is rounded, not an exact count.
Collect each block’s residual output before final normalization and plot mean vector norm over batch/positions. Independent zero-mean updates would make squared norms grow roughly linearly, giving square-root growth of norms; correlated updates and initialization can change this. Measure rather than assuming monotonic growth.
Post-norm computes a = LN1(x + attn(x)); y = LN2(a + mlp(a)). With zero branches it returns LN2(LN1(x)), generally different from x, because normalization remains on the residual path.
Exact GELU is $x(1+\operatorname{erf}(x/\sqrt{2}))/2$. A dense scalar sweep over [-6, 6] gives a maximum absolute difference of about $4.73\times10^{-4}$ from the tanh approximation. That is an activation error, not a logit bound: compare fixed-input logits after replacing every GELU, since later weights and layers can amplify or cancel it.

7. Training

Check your understanding

It applies a numerically stable log-softmax internally (subtracting the maximum, using log-sum-exp). Passing probabilities would apply softmax twice, or take the log of rounded probabilities and lose precision.
A uniform guess gives $\ln 50{,}257 \approx 10.8$. A much larger initial loss means the model is confidently wrong: initialization too large, a missing scale, or mismatched targets.
Overfitting: the model memorizes training text, which lowers training loss, while its predictions on unseen text get worse.
Teacher forcing always feeds the correct previous tokens. At generation time, the model feeds its own outputs, so one early mistake shifts it into contexts it never saw; loss also doesn’t measure qualities like repetition.
The optimizer state (AdamW moments and step count), the learning-rate schedule position, the random number generator states, and the data order position.

Stretch exercises

On each validation improvement, clone the state tensors or copy the state dict; save patience and best loss, then restore the best weights after stopping. Holding references to live tensors would overwrite the saved best state. Judge sample changes with fixed prompts/seeds and report best held-out loss; improvement in readability is not guaranteed.
Change one setting per run with the same data split and training budget. More dropout can reduce overfitting but slow fitting; larger width increases capacity and cost; 1e-2 can converge faster or destabilize training compared with 1e-3. Record best validation loss rather than assuming any setting wins.
With comparable data quality and enough optimization, more distinct text usually reduces memorization and the train/validation gap. Preserve document-level splits and report tokens processed as well as steps; a fixed step budget may undertrain the larger corpus.
Zero gradients once, sum each microbatch’s token losses divided by the total valid-token count across all microbatches, then clip and step once. Compare with the concatenated batch using dropout off and identical initial optimizer state. Unequal token counts require token weighting, and different reduction order can introduce small floating-point differences.

8. Generation

Check your understanding

Dividing logits by 0 is undefined, and the limit as temperature goes to 0 is the argmax. A separate branch computes the argmax directly, deterministically.
The mass before token 0 is 0 (keep), before token 1 is 0.4 (keep, since 0.4 < 0.5), before token 2 is 0.7 (drop). Tokens 0 and 1 remain, renormalized to 4/7 and 3/7.
So that a request’s samples depend only on its own seed, not on how many other requests used a shared generator before it. Reproducibility and isolation require it.
The prompt’s last position already produces the logits for the next token during prefill.

Stretch exercises

Decode incrementally and check the accumulated suffix, including stops spanning token boundaries and UTF-8 pieces. Hold back suffixes that could begin a stop string before streaming them. A stop may end inside one token’s text, so trim the decoded output at the match rather than relying only on token IDs.
For output count c_i, subtract presence_penalty * (c_i > 0) and frequency_penalty * c_i from logit i. Repetition penalty instead divides positive seen logits and multiplies negative ones by its factor. Compare repetition and held-out quality at matched settings; stronger penalties can suppress useful repeated words.
Keep four candidate histories, expand them using accumulated log-softmax scores, and retain the best four while tracking finished hypotheses separately. Compare summed logprob at fixed length or specify a length penalty. Beam search often finds higher-scoring sequences, but finite-width search need not beat greedy globally and probability need not predict readability.
Define a grammar, for example -?[0-9]+(\.[0-9]+)?, and track states for sign, integer, decimal point and fraction. For each state, allow a vocabulary token only if all its bytes follow valid transitions; allow EOS only in accepting states. A single context-independent “numeric token” flag cannot enforce the grammar.

9. Real weights

Check your understanding

The header is a JSON block at the start of the file, preceded by its length, listing every tensor’s name, dtype, shape and byte range. Reading it doesn’t touch the data section.
Pickle can execute arbitrary code during loading. A malicious checkpoint can run commands on your machine; safetensors contains only data.
A square matrix and its transpose have the same shape, so the copy succeeds; the outputs are wrong. Only comparing values with a reference reveals it.
With tied weights, the head is the embedding matrix, so storing it twice would be redundant. The loader must point the head at the embedding (or verify equality if both are present).
The output of the first layer (or the embedding), then each layer in turn: the first point where values diverge contains the bug.

Stretch exercises

Sum numel * dtype_size over stored tensors and compare with the tied parameter formula times 4 for an FP32 checkpoint. Additional bytes may come from duplicated tied weights or stored buffers, such as attention masks in some checkpoint formats. Report names/dtypes rather than assuming every GPT-2 snapshot stores the same extras.
Load the same weights in both dtypes, evaluate fixed token IDs, and report max-absolute and relative-L2 logit error. BF16 rounding can flip near-tied argmax choices; after the first flip, free-running texts no longer isolate precision error. Neither identical text nor a particular error magnitude is guaranteed.
The attention c_proj is square, so its shape check passes; reference layer/logit comparisons catch the wrong orientation. The MLP c_proj is rectangular, so skipping its transpose also fails shape checks. Generated text can degrade, but text inspection alone is an unreliable correctness test.
Export a matching config and Hugging Face names, transposing linear weights into GPT-2’s Conv1D convention while leaving embeddings/norms unchanged. Preserve biases, tied-head semantics and tanh GELU. Reload with transformers in eval mode and compare logits for several fixed lengths and batches.

10. GPU performance

Check your understanding

Kernel launches are asynchronous. Without synchronizing first, earlier queued work is still running and gets counted in your timing.
Profilers add overhead and change timing. Use the profiler to find where time goes, and separate clean runs to measure how much.
Each weight byte supports about 16 FLOPs (2 FLOPs per BF16 weight × 16 sequences ÷ 2 bytes). That’s far below the H100’s ridge point of about 295 FLOPs per byte, so it’s memory-bound.
Amdahl’s law: $1 / (0.7 + 0.3/3) = 1/0.8 = 1.25\times$.
DGX Spark’s CPU and GPU share one LPDDR5x memory with 273 GB/s of bandwidth. CPU-heavy work consumes bandwidth the memory-bound decode needs.

Stretch exercises

Unsynchronized wall time mostly measures enqueue/launch time and can miss the GPU computation. Synchronizing before and after gives end-to-end completed-operation time; CUDA events measure the interval on the GPU stream. Warm up first and use the same TF32 setting in every run.
For FP32 x+y, count two reads and one write: achieved bandwidth is 12*N/time bytes/s. Small vectors are launch-bound; large vectors approach a bandwidth plateau. Use repeated synchronized runs and distinguish cache-resident cases from DRAM traffic; the plateau’s fraction of peak is hardware-dependent.
Use $2N^3/t$ FLOP/s and approximately $N/3$ FLOPs/byte for square BF16 GEMM under a one-read-per-input, one-write-output traffic model. Small matrices underfill the GPU; large ones can approach compute peak. Whether any size reaches 70% depends on device, clocks and library kernels.
Inspect a GPU timeline, count kernels in one warmed-up decode step, and compare their occupied-time union with the step interval. Gaps identify time without kernel execution, although dependencies and transfers also cause gaps. Measure final latency in a separate unprofiled run.

11. CUDA

Check your understanding

Block sizes come in fixed sizes (multiples of 32, often 256), and $N$ is rarely a multiple. Rounding the grid up and having extra threads return keeps the code simple and the block size efficient.
CUDA doesn’t guarantee that all blocks run at the same time or in any order. A block waiting for another that hasn’t been scheduled can deadlock.
Addresses are 64 bytes apart, so the warp spans 32 × 64 = 2,048 bytes: 16 segments of 128 bytes. Only 32 × 4 = 128 bytes are used: 6.25% efficiency.
Each inner iteration loads two 4-byte floats for 2 FLOPs: 0.25 FLOPs per byte. The roofline predicts speed limited to 0.25 × bandwidth, a tiny fraction of peak compute.

Stretch exercises

Load/store four aligned FP32 elements per thread and handle the tail separately when N is not divisible by 4. Vectorization can reduce instruction overhead but does not reduce the three arrays’ total bytes; speedup can be small when scalar vector add already saturates bandwidth.
Compare equal numbers of reads/writes using CUDA events. A stride of 32 spreads a warp’s reads across separate memory regions and wastes transactions. 32*t % N repeats indices for many N values, so also use a bijective strided permutation to distinguish coalescing effects from cache reuse.
With consecutive lanes mapped to rows, row-major A reads and output stores become strided instead of contiguous, increasing memory transactions. The precise slowdown depends on broadcasts, caches and matrix size; verify the output before timing.
Load a 32×32 tile with contiguous lanes reading contiguous columns, synchronize, then swap tile/thread coordinates for contiguous output writes. Declare the shared tile as [32][33] to avoid transpose bank conflicts. Mask edge loads/stores while ensuring every thread reaches the barrier.

12. Fast matrix multiplication

Check your understanding

The first barrier ensures the whole tile is loaded before anyone computes with it. The second ensures everyone has finished computing before the next phase overwrites the tile. Without the second, fast threads overwrite data that slow threads are still reading.
Larger tiles need more shared memory and registers per block, so fewer blocks fit on an SM (lower occupancy), and edge tiles waste more work. Past a point, lost parallelism outweighs the reuse.
Column 5 is at word addresses $32i + 5$, all in bank 5: a 32-way conflict, serialized into 32 accesses. With row length 33 the addresses are $33i + 5$, in banks $(i + 5) \bmod 32$, all different: no conflict.
Each weight is used exactly once in a matrix-vector product, so there’s no reuse for tiling to exploit. It stays memory-bound; only reading fewer bytes (quantization) or batching helps.

Stretch exercises

For 256³ work, input loads are $2(256)^3/b$: 8,388,608, 4,194,304, 2,097,152 and 1,048,576 elements for tiles 4, 8, 16 and 32. Two 64×64 FP32 tiles need 32,768 bytes (32 KiB), before padding or double buffering. Compare with the per-block shared-memory limit; a 64×64 one-thread-per-output block would also exceed the usual thread limit.
Give each thread four accumulators, reuse two A and two B register values per inner step, and write its 2×2 output region. This reduces repeated shared-memory reads at the cost of more registers. Check irregular boundaries and measure occupancy as well as time.
A warp strides through one W row with consecutive lanes loading consecutive columns, accumulates in FP32, then reduces with warp shuffles and lets lane 0 store. Mask the K tail and count at least 2*N*K weight bytes for bandwidth. Vector reuse and caching affect the full traffic accounting.
Stage correctly aligned BF16 tiles, load WMMA fragments, accumulate in FP32 and store the result with the required layout. Validate against cuBLAS with dtype-appropriate tolerances. cuBLAS generally has better pipelining and tuning; multiples of 16 remove edge handling but do not guarantee competitive speed.

13. Reductions and numerics

Check your understanding

Softmax is invariant to adding a constant to all inputs: $e^{x_i - c} / \sum e^{x_j - c} = e^{x_i} / \sum e^{x_j}$.
BF16 keeps FP32’s 8 exponent bits (same range) but only 7 mantissa bits (about 3 significant decimal digits) instead of 23.
Adding a small number to a large BF16 sum loses the small number’s low bits, or all of it (swamping). FP32 accumulation keeps them; the result can be rounded to BF16 at the end.
The maximum error can be dominated by one harmless outlier, or hide a systematic small error everywhere. An aggregate (relative L2, mean) describes the typical error.
$-\infty$, the identity of max: it never changes the result.

Stretch exercises

Use Kahan’s update y = value - c; t = s + y; c = (t - s) - y; s = t, explicitly rounding every operation to BF16 for this experiment. With explicit round-to-nearest-even BF16 rounding after every operation, this recurrence returns 1,004 instead of the plain accumulator’s 32: an error of 4 relative to 1,000. The represented input is 0.10009765625, whose exact 10,000-term sum is 1,000.9765625: compensation does not undo input rounding or guarantee exact summation in such low precision. Compare errors against both reference sums.
Initialize active lanes from data and inactive lanes to negative infinity, then reduce with max. Padding with 0 makes every all-negative row incorrectly return 0. Test non-power-of-two widths as well as negative inputs.
Evaluate fixed histories and report max(abs(a-b)) and norm(a-b)/norm(a), with a defined convention for a zero reference norm. Also compare 50-step greedy continuations and record the first divergence. Precision, kernels and model determine the numbers; compare teacher-forced logits after divergence.
Load a row into registers, reduce its maximum, compute exponentials relative to that maximum, reduce their sum and normalize before storing. Use FP32 intermediates and mask inactive lanes. Ideal traffic is one read and one write per value; very long rows may spill registers or require a multi-block algorithm.

14. Triton

Check your understanding

One program instance handles one block of data with vector operations; Triton maps it onto many threads (a few warps). A CUDA thread handles individual elements and you manage the cooperation yourself.
Softmax is memory-bound, so time is proportional to memory passes. The fused kernel reads and writes each element once instead of about eight times.
Masked lanes must not affect the maximum. Filling them with $-\infty$ makes them neutral; 0 would be wrong for rows of negative values.
To verify the kernel on every input, to fall back when the kernel doesn’t support a shape or device, and to measure the kernel’s actual benefit.

Stretch exercises

Apply GELU to the FP32 accumulator before its final cast/store. Match the reference’s exact or tanh convention; applying GELU before rounding the matmul output can differ slightly from F.gelu(a @ b) with a low-precision intermediate. Compare with an appropriate reference and tolerance.
Autotune valid tile sizes, warp counts and stage counts keyed by matrix dimensions and dtype. Exclude compilation/tuning from steady-state timing, validate every selected configuration and compare with torch.matmul under identical settings. The winning configurations and speedups vary by GPU and shape.
Load both members of each rotary pair before writing either one, apply a*cos-b*sin and b*cos+a*sin, and use absolute positions. Respect split-half pairing, rotary dimension, strides and differing Q/K head counts. Compare with apply_rope on multiple positions and tails.
Reduce tile-local maxima/sums to obtain a row maximum and denominator, then reread and normalize the row. Alternatively update a running maximum/sum with exponential rescaling while scanning tiles, then make an output pass. Neither approach can finalize early tiles’ normalized outputs before the denominator is known.

15. FlashAttention

Check your understanding

Its terms are measured relative to a running maximum that may still change. Normalizing early would need rescaling again; dividing once by $\ell$ at the end is cheaper and exact.
Each tile’s softmax normalizes by that tile’s sum only. The true weights normalize by the sum over all tiles, so the per-tile results are wrongly weighted.
Intermediate memory (from $O(T^2)$ to $O(T)$) and memory traffic. The arithmetic is unchanged.
Decode has one query per sequence and head, too few programs to fill the GPU. Splitting the keys across programs and merging the partial results restores parallelism.

Stretch exercises

Reset peak allocation statistics for each run and compare peaks above the same baseline. One FP32 4,096² score matrix costs 64 MiB per batch/head; materializing both scores and probabilities costs roughly twice that. Online attention replaces these with tile-sized scratch and row accumulators, but actual savings depend on implementation and live tensors.
Let m = max(m_s) over splits, then merge l = sum(exp(m_s-m)*l_s) and a = sum(exp(m_s-m)*a_s). Return a/l. Treat empty splits as zero mass and return zero for all-empty rows; compare different partitions with the unsplit reference.
For a window containing w keys including the current position, valid keys satisfy q_pos-w < k_pos <= q_pos. Skip tiles whose keys all lie outside this range and mask the boundary tiles. Verify queries at the start of the sequence and after an existing prefix.
Reconstruct $P=\exp(S-\operatorname{LSE})$ tilewise using the original mask and score scale. Compute $dV=P^\top dO$, $dP=dO V^\top$, $dS=P\odot(dP-\sum_j P\odot dP)$, then $dQ=dS K/\sqrt d$ and $dK=dS^\top Q/\sqrt d$. Accumulate contributions across tiles, sum shared-head gradients for GQA and compare with autograd.

16. The KV cache

Check your understanding

A query is used only once, by its own position, at the step it’s computed. Keys and values are read by every later position.
Prefill: the prompt’s last position predicts the first new token.
With a 3-token prefix, the chunk’s queries are at positions 3 and 4 and keys at 0-4. A square triangle aligned at the top-left would let query 3 see only key 0 and query 4 only keys 0-1. The rule key_pos <= query_pos gives the correct rectangle.
The cache stores KV heads, which GQA reduces. Every query head still computes scores against the shared keys, so the number of scores depends on query heads.
Each decode step reads the whole cache once, and the cache grows with context while the weights don’t. At long contexts the KV bytes exceed the weight bytes.

Stretch exercises

Repeated concatenation copies the growing history: total copied elements scale as P*N + N*(N-1)/2 for prompt length P and N new tokens, per cached stream. Preallocated writes copy only the new entries. Measure synchronized generation at each length; the point where quadratic copies dominate depends on model and device.
Qwen3-0.6B’s BF16 cache at 32,768 tokens is $28\times2\times8\times32{,}768\times128\times2=3{,}758{,}096{,}384$ bytes (3.5 GiB). An ideal decode reads roughly this amount plus about 1.2 GB of BF16 weights, so KV is about three quarters of their combined traffic. Report actual stored weights, context length and measured reuse separately.
Store position p in physical slot p % w, but retain absolute positions for RoPE and the causal/window mask. Gather in logical order or pass each slot’s absolute key position; physical slot indices cannot stand in for positions. Protect keys still needed by earlier queries in a multi-token chunk.
Store FP8 values and scales, dequantize inside attention, and ensure a changing scale never reinterprets older entries encoded with a different scale. Fixed calibrated per-head scales or explicit rescaling are possible choices. Compare against BF16 on the same 1,000 teacher-forced positions and report errors by context length; free-running divergence is a separate result.

17. Qwen3

Check your understanding

A reshape assumes a layout. If the projection’s output features are ordered differently from how you split them into heads (for example, head-major versus interleaved), the shapes match but heads get the wrong features.
Project with k_proj, reshape into KV heads, apply QK-norm (RMSNorm per head), then RoPE at its absolute position. The rotated, normalized key is what’s cached.
Rotations compose: rotating $q$ by $m\theta$ and $k$ by $n\theta$ gives a dot product that depends only on $(m - n)\theta$. Shifting both positions by the same amount leaves the difference unchanged.
SwiGLU computes a gate and a value from the input with two matrices, multiplies them elementwise ($\operatorname{SiLU}(xW_g) \odot xW_u$), then projects down with a third. GPT-2’s MLP has one up projection and one down.
The architecture is the computation, independent of the weight values. Identical random weights in both implementations exercise every operation, and any mismatch is a bug in the computation. Loading real weights then only adds name mapping, which has its own checks.

Stretch exercises

Each rotary pair is multiplied by an orthogonal 2×2 rotation, so its squared norm is unchanged. At equal positions, $R_p^\top R_p=I$, making the rotated Q/K dot product equal to the unrotated one. Verify with floating-point tolerances across several positions.
Rotate features (0,1), (2,3), ... instead of split-half pairs. Using unchanged Qwen3 weights applies different positional geometry and changes nonzero-position logits; compare fixed-input max-absolute and relative-L2 errors. An appropriate Q/K row permutation can reconcile conventions, so pairing and weight layout must change together.
Linear extension by factor f uses position/f, equivalently frequencies divided by f, multiplying wavelengths by f. Multiplying positions by f instead shortens wavelengths. YaRN scales slow frequencies while preserving fast ones and blending between them; match the model’s attention scaling too.
For each Q/K head, reduce mean square in FP32, apply its learned RMSNorm weights, then rotate paired components and store once. Handle different Q and KV head counts and the checkpoint’s epsilon/pairing. Keep a reference backend and compare fixed-input logits before benchmarking the fused path.

18. Engine v1

Check your understanding

Generated text can look fine with subtly wrong computation, and once one token differs everything after differs. Comparing layer outputs on fixed inputs finds where the computation first diverges.
The prompt adds a KV cache proportional to its length, plus activations and logits proportional to it during prefill (all 4,096 × 151,936 logits in FP32 is 2.5 GB). Weights fitting says nothing about those.
Prefill processes all prompt tokens into the cache but emits one token; later steps process one token each. The last emitted token is never processed into the cache unless generation continues.
to_empty allocates fresh storage for every parameter, including the head, which breaks the shared storage. The head must be pointed back at the embedding.
<|im_end|> (end of the assistant turn) and <|endoftext|> (end of document). With only <|endoftext|>, the model ends its turn and continues by writing an imaginary next turn.

Stretch exercises

Slice final normalized hidden states to the last position before the head. With vocabulary 151,936 and FP32 logits, avoiding 4,095 rows saves $4{,}095\times151{,}936\times4=2{,}488{,}711{,}680$ bytes, about 2.32 GiB, plus unnecessary head computation. This does not remove prefill activations or KV.
Register corresponding layer hooks, detach outputs, align their layouts/dtypes, and report the earliest layer outside absolute/relative tolerances. Compare embeddings first and use identical token IDs, positions and eval settings. Remove hooks in finally and avoid retaining all GPU outputs when memory is tight.
Pass the thinking option through the checkpoint’s chat template and record correctness and generated-token counts on the same five questions. Thinking mode often uses more tokens for intermediate reasoning, but quality improvements are task-dependent. State whether the budget includes reasoning tokens and inspect the rendered prompts.
Parse vocabulary, ranked merges, added tokens, byte mapping and the exact Unicode-aware pre-tokenizer. Match normalization and added-token rules before BPE. Compare all 1,000 lines with AutoTokenizer, including whitespace, Unicode and reserved tokens; Chapter 35 develops this complete boundary.

19. Fast decode

Check your understanding

int(token) makes the CPU wait until the GPU finishes, emptying the GPU’s queue. The next step’s kernels then start only as fast as the CPU can launch them, and the GPU idles between them. Without the sync, the CPU queues work ahead.
A graph replays recorded kernels; it has no CPU logic. The if would need the value on the CPU, a sync that isn’t allowed during capture, and its branch would be frozen at capture time.
Unwritten slots sit at positions greater than every live query’s position, so the causal rule key_pos <= query_pos masks them.
CPU operations run synchronously, so there’s no asynchronous queue to keep full and nothing for sync removal to gain.
The graph keeps reading the old tensor’s address, so every replay feeds the same stale token. In-place updates keep the address the graph recorded.

Stretch exercises

Typical v1 synchronizations include GPU scalar extraction with int, .item() or .tolist(), CPU copies for output, and host-dependent stop checks. The fast loop defers these until a batched result transfer. Record the actual warning sites; sync-debug warnings do not exhaustively identify every synchronization.
Concatenate Q/K/V weights along their output dimension and gate/up likewise, then split the outputs using their original widths. This replaces five projection matmuls with two per layer, saving three matmul launches before other fusions. Actual profiler kernel counts also depend on the library implementations.
Use the fused residual-add/RMSNorm operation while preserving both the updated residual and normalized input required by later branches. Verify layer outputs and logits, including epsilon and dtype, then time warmed graph replay. Gains depend on how much of decode is launch-bound.
Right-pad to the bucket so real causal queries cannot attend to padding; keep real positions unchanged and give padding isolated/scratch cache slots. Use a validity mask and select logits at the last real position, not the bucket’s last row. Padding positions alone do not protect cache writes or output selection.
Use fixed-shaped sort/top-k tensors and mask logits in place; keep the token crossing the top-p threshold by testing cumulative mass before each token. Sample with fixed-shaped device operations and per-request RNG. Test empirical frequencies against the reference rather than expecting identical RNG draws across different algorithms.

20. Quantization

Check your understanding

Decode is memory-bound, so time is proportional to bytes read. Large-batch prefill is compute-bound, and its arithmetic doesn’t shrink with fewer bits (unless the math itself runs at lower precision).
Each scale covers fewer values, so a large value raises the step size for fewer neighbors. The cost is more scales: metadata bytes and work in the kernel.
Clipping saturates rare large weights (more weight error) but gives all others a finer grid. If the clipped weights multiply inputs that are usually small, the output improves.
It materializes the full-precision weight matrix and reads it again, so as many bytes cross the memory bus as without quantization (more, counting the dequantization pass).
Their errors are amplified: the head’s errors land directly on the logits, and a router’s small error can flip a discrete expert choice. They’re also small, so quantizing them saves little.

Stretch exercises

For b-bit codes and one s-bit scale per group of g weights, effective bits/weight are approximately b + s/g, plus any zeros, padding or other metadata. Smaller groups generally reduce held-out output error but increase metadata. Plot error on held-out activations, not only weight reconstruction error.
Dispatch CUDA calls to a kernel that unpacks and dequantizes tiles in registers while multiplying BF16 activations. Avoid constructing a full BF16 weight matrix, keep graph buffers stable and retain a reference path. Validate logits and compare memory and warmed decode speed; a small launch-bound model need not realize a fourfold gain.
Clamp activation statistics away from zero, search alpha on calibration data, quantize W*diag(s) and evaluate with x/s. Select by output error, then evaluate on held-out inputs. Fold inverse scaling into a compatible preceding operation only when every affected consumer is adjusted consistently.
Build a damped activation Hessian, factor its inverse, and for column i quantize w_i, form e_i = (w_i-q_i)/U_ii, then subtract e_i*U_i,j from unprocessed columns. Handle poorly observed columns and group boundaries. Compare held-out output error with RTN at matched bits/group size; Chapter 39 gives the full implementation.

21. Fine-tuning

Check your understanding

Under the causal mask, only the last token has attended to the whole input. The first token has seen only itself.
The last prompt token’s position: its target is the first response token.
Masking by ID would also mask the real EOS that ends the response, so the model would never learn to stop.
Training needs gradients and optimizer state for every trained parameter (and an FP32 master copy in mixed precision): about 18 bytes per parameter against 2, plus activations saved for the backward pass.
Fine-tuning on 935 examples teaches a pattern (the answer format and when to stop). Facts need far more text than that; a pretrained model already has them.

Stretch exercises

Prompt-inclusive loss spends updates learning prompt tokens as well as responses; response-only loss targets the desired completion behavior directly. Under the same budget it often improves response learning efficiency, but there is no guaranteed winner. Evaluate both with the same response-only held-out mask and generation prompts.
Use the last non-padding hidden state for classification, keep one fixed train/validation/test split, and count only parameters with gradients enabled. Head-only training is cheapest; unfreezing the last block or all layers adds adaptability, memory cost and overfitting risk. Select hyperparameters on validation data before reporting test accuracy.
Concatenate examples, restart their position IDs, and allow attention only within each example to earlier/equal positions. Mask cross-example next-token targets and retain response-only loss boundaries. Compare token-weighted loss with unpacked examples; packing helps most when padding waste was large.
Let $\Delta=\log\pi_\theta(y_c|x)-\log\pi_\theta(y_r|x)-\log\pi_{ref}(y_c|x)+\log\pi_{ref}(y_r|x)$. Minimize $-\log\sigma(\beta\Delta)$ using response-token logprob sums and a frozen reference. Mask prompt/padding tokens and compare held-out preference accuracy and ordinary response quality.

22. LoRA

Check your understanding

The adapter’s contribution is $\frac{\alpha}{r} BAx$, which is zero when $B = 0$.
$\partial L/\partial A$ contains $B^\top$, which is zero. $B$’s gradient isn’t zero, so after one step $B \ne 0$ and $A$ starts receiving gradients.
Merging costs nothing at inference (an ordinary weight matrix). Separate adapters let one base model serve many fine-tunes, swapped per request, each a few MB.
The base is stored in 4 bits (about a quarter of the memory), but each forward and backward pass must dequantize it, which costs time.
The adapter learned a correction relative to specific weights. On different weights, the same correction no longer corresponds to the intended change.

Stretch exercises

Each adapted linear layer adds r*(d_in+d_out) parameters, so parameter count grows linearly with rank. Run identical data/step budgets and multiple seeds; plot held-out loss against the actual count. The rank where gains flatten is task- and layer-selection-dependent.
LoRA reduces trainable gradients and optimizer state substantially, while still needing base weights and activations for adapter training. Compare held-out loss, peak allocation and steady-state step time with matched data and steps. Lower memory is expected; equal quality or faster steps is not guaranteed.
Compute the shared base projection once, gather each row’s A/B matrices and add its scaled adapter correction. Compare with separate row runs and cover a no-adapter row. Preserve each adapter’s rank/scale and prevent one row’s weights from affecting another.
Initialize m to the base weight’s norm along the exercise’s column convention and compute the normalized direction from W + BA, with a safe denominator. Count m’s extra parameters when choosing a rank for the comparison. Check that zero adapter correction reproduces W; translate the normalization axis if using [out, in] storage with a different convention.

23. Inside the model

Check your understanding

The residual stream at every layer lives in the same space (each layer adds to it), and the head was trained to read that space. Early layers aren’t trained to be decoded, so the lens is approximate there.
The direction is a difference of means. Anything else that differs between the groups (topic, length) also ends up in the difference and gets steered or removed along with the property.
Adding moves every activation along the direction by a fixed amount. Projecting out sets the component along the direction to exactly zero, whatever it was.
The residual stream is a sum of what the embedding and each layer write. If none of them can write along $d$, the stream never has a component along $d$.
Measure the target behavior on held-out prompts, and the KL divergence from the original and held-out loss on unrelated data. Small changes elsewhere show a targeted edit.

Stretch exercises

Learn the direction using training examples at each layer, apply the same 1.0× mean-gap convention and evaluate all-caps probability on held-out prompts. Plot the baseline as well as the steered probability. The best layer is empirical and can change with model or prompt distribution.
At every layer, apply the final norm and head to the last prompt position, then rank the exact token ID for Paris. Report the first top-five entry and its probability; it may leave and reenter. The layer depends on the checkpoint and tokenizer, so it cannot be inferred from architecture alone.
Select a layer on held-out calibration prompts, then generate from identical test prompts with lambda 0, 4 and 8. Larger additions usually strengthen the selected style but can distort content or fluency. Record the direction’s normalization and compare unrelated-task loss/KL as well as the five answers.
For unit direction d, project each output writer with W' = (I-dd.T)W and b' = (I-dd.T)b; project embedding rows similarly. The runtime comparison must project those same writes, not just one selected layer. Preserve head tying consistently: if editing a tied embedding also edits the head, apply the equivalent head change in the hooked reference. Save/reload and compare fixed-input logits.

24. Batching

Check your understanding

All requests share the weights, so one read serves the whole batch. Each request has its own KV cache, so cache reads grow with the batch.
Finished requests’ slots idle until the whole batch finishes, and new requests wait for the slowest member. Iteration-level scheduling refills slots every step.
Decode latency is what streaming users feel; a large prefill in the same step would stall every running request. Serving decode first and filling the remaining budget with prefill chunks keeps time per token steady.
Each row has its own positions entry and its own cache row (slot). Attention masks by key_pos <= query_pos per row, so rows at different lengths need no padding.
So that a request’s samples don’t depend on which other requests share its batch, keeping results reproducible and isolated.

Stretch exercises

Throughput generally rises as more requests reuse each weight read, then plateaus as compute, KV traffic or capacity becomes limiting. Mean TPOT can grow even while aggregate throughput improves. Use the same lengths and sampling settings and report both curves; the saturation slot count is hardware-dependent.
Draw interarrival times from an exponential distribution with mean 1/rate, and measure TTFT from the scheduled arrival. Smaller token budgets limit prefill interference but can increase queueing/prefill steps; larger budgets trade better utilization for longer decode stalls. Report p50/p99 at the same offered load and include failures.
Give each request an output queue, submit work to one owning engine loop, and stream incremental text as SSE events. Validate/admit before sending headers and abort on disconnect. Keep tokenization and detokenization from blocking decode; Chapter 36 develops the process-separated server.
Save prompt/output IDs, progress and RNG state, retire the victim’s slot, then rebuild its KV on readmission without sampling its already committed tokens again. Compare with solo execution using per-request generators and fixed settings. Bitwise equality may also require shape-independent numerics or a deterministic reference backend.

25. Paged attention

Check your understanding

The output length isn’t known until generation stops, so exact allocation is impossible; reserving the maximum wastes memory, and growing contiguous buffers means copies and fragmentation.
In physical block table[p // b], at offset p % b.
A block’s keys and values depend on every earlier token. Two blocks with the same tokens but different histories have different K/V.
When a sequence is about to write into a block whose reference count is above 1. Usually the partially filled last block of a shared prompt.
The prompt’s last position must run through the model to produce the logits for the first output token.

Stretch exercises

Utilization is sum(lengths)/(block_size*sum(ceil(length/block_size))) for nonempty requests without sharing. Smaller blocks reduce tail waste; larger blocks reduce block-table and allocator overhead. Plot the same length set for every size and report metadata separately; a utilization peak alone does not choose the fastest size.
Reserve every chunk before model execution, release evictable zero-reference prefixes first, and preempt a live request if capacity is still insufficient. Keep block reference counts and copy-on-write correct during adoption and writes. Verify request outputs and that every block is reclaimed after completion/abort.
Prefill once, fork references to shared blocks, and give each child independent sampling state. A child writing into a shared partial block must copy it first. Compare each branch with a solo request using its own seed and verify freeing one child does not invalidate the others.
Map each query’s absolute position through its request’s block table and stream paged K/V using online softmax. Include the newly written chunk only up to each query’s causal position and keep ragged requests isolated. Compare against gathered dense attention across partial blocks and nonzero prefixes.

26. Speculative decoding

Check your understanding

At batch 1, a decode step is memory-bound: reading the weights dominates. Five tokens read the same weights once, and the extra arithmetic is small.
Through acceptance, $y$ is output with probability $q(y)\min(1, p(y)/q(y)) = \min(p(y), q(y))$. A rejection happens with probability $\sum_z \max(0, p(z) - q(z))$ and then outputs $y$ with probability $\max(0, p(y)-q(y))$ divided by that same sum. The total is $\min(p, q) + \max(0, p - q) = p(y)$.
If the first proposal is rejected, the target’s own choice (or the residual sample) replaces it; if all are accepted, the bonus token is added.
The target and draft wrote K/V for proposals that may have been rejected. Both are truncated to the committed tokens except the newest one, which hasn’t been fed yet: base length + accepted + 1.
At large batches decode is compute-bound, so verifying $k+1$ tokens costs about $k+1$ times as much, and rejected tokens are wasted compute.

Stretch exercises

Measure acceptance, draft time and verification time for each k. With approximately constant conditional acceptance alpha, expected emitted tokens per round are $E=(1-\alpha^{k+1})/(1-\alpha)$, or k+1 at alpha=1; predicted speedup is E*t_decode/(t_draft(k)+t_verify(k)). Observed acceptance varies by depth and prompt, so the formula is a model, not a guaranteed speedup.
Search for an earlier occurrence of the trailing three-token sequence that has following tokens available, and propose up to k of those tokens. Treat each proposal as deterministic and verify with the target’s greedy or exact sampling rule. Rewriting often reuses source spans; empty matches or changed wording reduce acceptance.
Estimate recent acceptance and draft/verify cost for candidate k values, then choose the largest expected tokens-per-second ratio, including k=0. Smooth estimates and periodically probe alternatives. Adaptation should respond to both acceptance and load, since longer verification can cost more than it saves.
Flatten the 14 draft nodes with depth-based absolute positions and mask each node to the shared prefix and its ancestors, including itself. Match target argmax choices along a single root-to-leaf path, commit only that path and obtain a correction/bonus from its endpoint. Keep sibling KV separate and compare with ordinary greedy decoding.

27. Mixture of experts

Check your understanding

Each token only runs, and so only reads, its active experts. But any token may choose any expert, so all must be resident.
Routing is a discrete choice. A rounding difference between two nearly tied experts can swap them and change the output completely.
Each expert’s tokens become contiguous rows, so one grouped GEMM can process all experts with dense matrix multiplications and no per-expert gather.
With $B$ tokens each choosing $k$ of $E$ experts, the number of distinct experts touched, $E(1 - (1 - k/E)^B)$, approaches $E$ quickly; by batch 32, Qwen3-30B-A3B touches 112 of 128 experts per layer.
An unbalanced router overloads a few experts and wastes the rest during training. At inference, the busiest expert sets the step time, and in expert-parallel serving the busiest GPU does.

Stretch exercises

Count top-k routed assignments per expert and layer on representative text, then compare dispersion, max/mean load or entropy. Each token contributes k assignments, so totals must match tokens*k. There is no universal early-versus-late balance pattern; use trained routers and report the observed distribution.
Group routed rows by expert, build prefix offsets and launch output tiles for each nonempty expert range. Use the matching expert weights, then scatter and combine outputs with routing weights. Handle repeated token destinations and empty experts; validate against per-expert reference matmuls before timing.
Quantize each expert’s matrices with local groups/scales while leaving the router and attention in BF16. Expert storage falls toward one quarter of BF16 plus metadata, but total model reduction depends on the expert fraction. Compare fixed-input logits and routing choices; later layers’ routes can still change through altered hidden states.
Keep host weights pinned, map expert IDs to resident GPU slots and copy misses before their kernels consume the slots. Do not evict an expert with in-flight users. Record transferred bytes, cache hits and stall time over the conversation; latency can be PCIe-bound despite a smaller GPU footprint.

28. Linear attention

Check your understanding

Without the exponential, $\sum_i (q_t \cdot k_i) v_i = q_t^\top \sum_i k_i v_i^\top$: the sum over the past can be accumulated once into a matrix and updated per token.
A query reads $\sum_i (k_i \cdot k_j) v_i$: every stored value leaks in proportionally to its key’s similarity. At most $d_k$ keys can be mutually orthogonal, which limits clean storage to $d_k$ associations.
The error $v_t - S^\top k_t$, scaled by $\beta$. With $\beta = 1$ and a unit key, after the update $S^\top k_t = v_t$ exactly.
Within a chunk, each position’s write depends on what earlier positions in the same chunk wrote (the prediction uses the updated state). That dependency is a unit lower-triangular linear system.
Every token’s update mixes into the same matrix; there’s no per-position entry to forget. Engines save snapshots of the state and restore one instead.

Stretch exercises

An untouched contribution decays as $0.95^n$, with a half-life of about 13.5 steps and about 0.006 of its initial amplitude after 100 steps. In an isolated recall example, error relative to the original value grows roughly as 1-0.95**n. Later corrective writes and key interference change this simple relationship.
Small chunks incur many sequential chunk transitions; large chunks increase local triangular work and temporary storage. A dense C×C triangular solve with C right-hand sides has cubic work per chunk, giving roughly quadratic per-position cost for that component. Time the proposed sizes at T=4,096; the best C depends on dimensions, implementation and device.
Fuse decay, old-state prediction, gated correction, state update and output read for each sequence/head, keeping the state tile on chip when it fits. Use FP32 state arithmetic as in the reference and avoid reading partially overwritten state for the prediction. Compare multi-step state and outputs; large states can cause register spills.
Remove the prediction-subtraction correction and use $S_t=\alpha_t S_{t-1}+k_t v_t^\top$ with readout $q_t^\top S_t$. Within a chunk, write contributions are weighted by products of intervening alpha values and the entering state by the full prefix product. This is the simplified selective recurrence in the exercise; removing correction does not by itself implement every component of a full Mamba-2 block.

29. Sparse attention

Check your understanding

To find the top $k$ scores it must compute all of them, which reads every key: the expensive part.
Blocks reduce the number of selection decisions and index keys by the block size, and the selected keys are contiguous in memory, which makes reads efficient.
The partially filled block containing the query’s own position. It’s always readable so recent tokens are never lost to selection.
All keys must stay stored, because a future query may select any of them. Only the reads per step shrink.
Selection is discrete: a tiny difference can pick a different block and change the output a lot. If the two paths broke ties differently, cached decoding wouldn’t reproduce the full forward.

Stretch exercises

Union the first four valid keys with the selected keys and remove duplicates before softmax. For flat scores, most probability mass may still lie outside the sparse set; for peaked scores, retaining the important peaks gives much smaller error. Ordinary retained “sink” keys here are different from Chapter 32’s learned denominator-only sink.
Maintain absolute positions and RoPE while mapping storage into position % w. Mask by the logical window, not by physical slot order, and prevent chunk writes from overwriting keys earlier queries still need. Compare multiple wraps and chunk sizes against the full-cache windowed reference.
Return selected block IDs, gather each block’s K/V and apply causal/token-valid masks within boundary blocks. Remove duplicates and define a sentinel for fewer-than-budget selections. Values should match the dense mask path within tolerance; actual sparse reads, rather than just a sparse mask, provide the traffic saving.
Iterate only selected blocks, load each through its physical mapping, and update the running maximum, denominator and weighted value sum. Apply causal and partial-block masks and handle empty selections. Compare with the explicit gathered reference and measure bytes/time as selection budget varies.

30. Capstone

Check your understanding

Gated DeltaNet (28), gated sparse attention with QSA and partial RoPE (29, 17), MoE with a shared expert (27), loading and quantization (9, 18, 20). New: hyper-connections and the n-gram memory.
Read weights are per stream and per feature (four 2,560-vectors), in (0, 1) through a sigmoid. Write weights are one scalar per stream, in (0, 2) through $2\sigma(\cdot)$.
Each head is an independent hash; distinct prime sizes make collisions in different heads independent. Two n-grams sharing a row in one head almost never share rows in all eight.
Each token reads 16 rows of 160 BF16 values, about 5 KB. Disk and page-cache reads of that size are small next to the gigabytes of weights read per token.
Keeping the old state is free, so a rejected speculation is undone by restoring it. A recurrent state can’t be truncated.
The dense parameters in BF16 (attention, linear attention, shared experts and LM head): about 8.5 GB of the 9.7 GB per token.

Stretch exercises

Snapshot paged KV progress and recurrent state before verification; after rejection restore both and rerun only the committed tokens to reconstruct the correct state. Synchronize the draft’s history/cache too. Compare complete greedy outputs with the nonspeculative path; a recurrent state cannot be fixed by KV truncation alone.
INT8 roughly halves the dense BF16 weight bytes, before scales, while leaving already-quantized routed experts unchanged. In the chapter’s estimate, halving 8.5 GB of a 9.7 GB/token weight read leaves about 5.45 GB/token, suggesting at most roughly 1.8× improvement if that traffic alone dominates. Validate fixed-input logits and measure actual latency; quantization kernels and non-weight work reduce the gain.
Inspect mtp. tensors/configuration, match the checkpoint’s embedding/hidden fusion, normalization and transformer module, and validate its forward computation before drafting. Use the target’s embedding/head where specified and maintain separate draft state. Report acceptance by draft depth and total latency; matching prefixes alone does not establish a correct MTP implementation.
Keep quantized expert codes/scales in pinned memory, stage only the routed experts and synchronize copy completion before use. Exact next-layer routes require its input; overlap using a verified pre-gating mechanism or speculative prefetch with a fallback for misses. Compare transferred bytes divided by achieved PCIe bandwidth with measured token latency.
Assign each request its own recurrent state and KV/block metadata, gather recurrent states into batches and scatter updated states back to their owners. Attention remains ragged and must isolate histories. Test joining, finishing, preemption and reordering against solo runs so slot reuse cannot mix recurrent states.

31. Engine v2

Check your understanding

Each computes the next n uncomputed positions of one request. Prompt/output boundaries change whether logits are sampled, not the scheduler’s work representation.
Attention needs sequence boundaries, lengths, block tables and slots; linear layers and tokenwise norms just see one matrix. Combining all prompt chunks gives larger matmuls and amortizes launches/weight reads.
The last prompt position must be computed to obtain the first output token’s logits. KV alone does not contain those logits.
Head-first insertion would make early prefix blocks older in the free queue and evict them sooner. Tail-first preserves the portion most likely to be reused.
Images and adapters change activations even when IDs match. Cryptographic hashes prevent deliberately crafted collisions from aliasing another request’s KV; all activation-changing inputs must enter the identity.
Preemption proves the pool is already under pressure. Admitting more work immediately can undo the reclaimed capacity and cause repeated eviction without progress.
Swap is likely cheaper for a modest context: moving a few KiB per cached token can cost far less than running billions of parameters over its history again. Compare context size, transfer bandwidth and recompute FLOPs on the target hardware; it is not universal.

Stretch exercises

Count adopted cached tokens, newly computed tokens and physically evicted cached tokens with clearly defined units. Multi-turn reuse improves when the pool preserves earlier full blocks, but the last position still needs logits and new text needs computation. Plot aggregate hit rate against pool size and avoid double-counting lookup hits that were not adopted.
Swap wins when host round-trip transfer plus bookkeeping costs less than recomputing the evicted history, especially when the history has little surviving prefix reuse. Prefix caching can make recomputation mostly a hit, reducing swap’s advantage. Force the same pressure/workload and compare total runtime and p99 TPOT with caching on and off.
Rank waiting requests by reusable prefix length, then compare hit rate, TTFT and throughput with FCFS. Popular-prefix traffic can indefinitely delay requests with short/no hits. Add age or bounded bypasses and measure worst wait rather than judging the scheduler only by hit rate.
Store token-block paths in a radix tree, split compressed edges at divergence, track live references and evict unreferenced leaves in LRU order. Copy a matched partial block before extending it. Full-block hash chains already share exact prefixes; any extra hit-rate benefit should come from partial matches or changed eviction behavior, not from the tree representation alone.

32. Attention kernels

Check your understanding

One decode query reads a large K/V history for relatively little arithmetic. Its memory lower bound is total required KV bytes across layers/requests divided by achievable bandwidth, before launch and other work.
Query heads sharing a KV head reuse the same loaded keys/values. Several heads become matrix rows, increasing reuse and offering a useful matmul tile even with one token per request.
Four requests give too few independent programs to occupy the GPU; partitioning each history supplies parallel work. Sixty-four requests may already supply enough programs, making split/merge overhead unnecessary.
Split s contributes partition sum exp(L_s) to the total softmax denominator, so its normalized output receives weight exp(L_s - logsumexp(L)). Partition sums and weighted numerators add associatively; numerical roundoff can still depend on order.
An empty split has zero probability mass, whose log is negative infinity. If every split is empty, subtracting negative infinity from negative infinity produces NaN; padding rows must instead return zero.
Different tokens and heads have different ranges. Local scales spend more of the quantizer’s levels on each range, reducing error at the cost of scale storage and reads.

Stretch exercises

Autotune only valid configurations using representative ragged lengths, head groups and hardware. Count useful KV bytes and compare achieved bandwidth with both attainable bandwidth and the predicted time floor. Exclude tuning/compilation, validate each configuration and include enough splits to occupy the GPU at small batches.
Compute shared-prefix attention and private-suffix attention separately and merge their outputs using log-sum-exp weights. Ideally the 4,000-token shared prefix is loaded once per query tile instead of independently for 64 requests, though real reuse is limited by tiling/cache capacity. Distinct queries still require distinct scores; measure traffic as well as latency.
Use the same held-out token sequence and teacher-forced loss with all cache formats, specifying scales, group axes and any unquantized residual region. INT8/FP8 often stay closer to BF16 than 4-bit caches, but actual perplexity changes are model-dependent. Per-channel key scaling addresses outliers; compare downstream behavior too.
For sink logit s, initialize the running maximum to s, denominator to 1 and value accumulator to zero, then process the real keys normally. This gives $O=\sum_j e^{z_j}v_j/(e^s+\sum_j e^{z_j})$. For split attention, include the sink exactly once in the final denominator, not once per split; test against dense attention with an appended zero-valued sink.

33. Overhead

Check your understanding

The smaller model’s memory work is short enough that many kernel launches and host decisions dominate. Reading the larger model’s weights takes much longer, so the same launch cost is a smaller fraction.
It saves launches, repeated input reads and separate matmul setup, while making one larger operation. The concatenated weight storage replaces the originals rather than duplicating them.
Data-dependent Python grouping and dynamically sized tensors/synchronizations cannot be replayed as a fixed graph. Device-side alignment uses fixed-capacity buffers and masked padded expert tiles instead.
Dummy rows read zero-length contexts and write only scratch slots. Sequence metadata isolates requests, so they cannot write real KV or join another request’s attention.
Captured kernels replay the bucket’s fixed program shape. Changing the host count underneath it would mismatch the captured metadata and actual padded rows.
Step s+1 may already have launched; its speculative extra token is discarded. Finishing frees the request’s references after the work is ordered, without exposing its placeholder KV as a reusable prefix.
The placeholder is not the actual token, so its hash does not identify the activation history. Only known committed IDs can be published under prefix keys.

Stretch exercises

Fusion lowers kernel count and repeated reads; graphs lower host launch gaps while replaying the captured kernels. Profile the same warmed-up step before/after and report occupied GPU time and idle gaps, then measure latency without profiling. Graph capture alone does not imply fewer numerical kernels.
Keep fixed-capacity metadata buffers keyed by stable request slots and update membership, lengths, token IDs and affected block-table rows in place. Lengths/positions of active decode requests still advance every step even when membership is unchanged. Compare host time at 256 requests and check that retired rows cannot retain live references.
Capture fixed-shape norm/projection/MLP pieces into token-count buckets, with eager attention between pieces. Keep interface buffers/addresses stable and isolate padding rows’ metadata and cache writes. Benchmark TTFT and decode latency under mixed load; extra boundaries and padding can offset launch savings.
Rotate normalized K with its absolute position, then write K and V directly into their assigned paged slots. Load each rotary pair before storing, respect head/stride layouts and mask padding/scratch rows. Compare pool contents and subsequent attention with apply_rope followed by write_kv.

34. Sampling and structured output

Check your understanding

Raw model logprobs describe the model distribution; temperature and truncation describe a sampling policy. Changing the policy should not change the reported underlying model score.
Repetition penalty rescales logits of seen IDs with sign-aware multiplication/division; frequency penalty subtracts a term proportional to output counts. Avoid indiscriminately penalizing source vocabulary in summarization; frequent source words may be exactly what the answer needs.
Independent exponential waiting times with rates p_i have token i as their minimum with probability p_i / sum(p). The race uses fixed-shaped device tensor operations and no host-selected index.
A shared generator advances when any neighbouring request samples, making a request’s output depend on batch membership. A per-request generator isolates its random stream, though numerical batch effects remain possible.
A prefix hit supplies KV, not logits for each prompt position. Prompt logprobs require recomputing those positions (unless separately cached).
Four normal requests arrive together; children adopt already published full prompt blocks of the first child. The final partial block, or last full block held back for logits, still runs for each child.
Token pieces can split Unicode characters and include arbitrary byte sequences. A character-level state cannot consistently advance across such token boundaries.
A fixed supported schema can bound structure/recursion and compile to a regular language. Arbitrarily nested JSON needs unbounded stack state and cannot be recognized by an ordinary finite automaton.

Stretch exercises

Zero a slot’s count row on admission, update counts for each committed output token and clear it on retirement/reuse. Prompt counts, if needed for repetition penalty, should be tracked according to that penalty’s semantics. INT32 storage costs 4*max_num_seqs*vocab bytes, so measure device memory alongside reduced host rebuilding time.
Store bit i in word i // 32, extract it with a shift/mask and replace forbidden logits with negative infinity. A packed mask uses about one eighth of a one-byte boolean mask’s memory, before padding. Validate vocab tails and signed bit 31; measure expansion/application time as well as mask transfer cost.
Combine byte-level JSON lexical states with a stack recording object/array context and expected separators/values. Precompute stack-independent transitions and simulate full token byte strings for stack-sensitive candidates. Handle strings, escapes, numbers and EOS only in accepting states, while bounding output length/nesting for resource control.
Find a forced byte continuation, then use a tokenizer-compatible continuation that preserves already committed token boundaries. Append forced tokens and update guide/text state, but process those tokens through the model before the next unconstrained choice so KV stays correct. This saves sampling/guide steps and can batch forced-position computation; reporting skipped model positions would misstate the work.

35. The text boundary

Check your understanding

The printable mapping makes every byte, including whitespace/control bytes, representable in JSON vocabulary strings without losing reversibility.
Added tokens are literal reserved boundaries. Normalization or regex splitting first could alter or split them so that their trained IDs are never emitted.
The regex determines which substrings BPE can merge. Different splits produce different IDs from the same merge table and therefore different model inputs.
One token can end halfway through a UTF-8 character. The incremental decoder buffers unfinished bytes until later pieces complete the character, instead of replacing each incomplete piece.
It holds back the trailing ST. It releases it when subsequent text disproves the stop prefix, or at normal finish; if OP follows, the completed stop is removed unless inclusion was requested.
Templates come from outside the server’s code and run on request data. A sandbox restricts object access and execution; it does not remove the need to bound rendering work.
A supported constraint masks every continuation outside the allowed tool-call syntax/name/schema. It guarantees membership in that language, not that a tool is appropriate or its arguments are factually correct.

Stretch exercises

Compare exact IDs on mixed whitespace, Unicode, special tokens and long strings, then inspect the first divergent substring. Differences can come from normalization, added-token rules, regex boundaries, merge ranking or byte mapping. Match the checkpoint’s supported tokenizer semantics rather than changing outputs to fit only the sampled corpus.
Maintain an incremental parser across chunks, emit the completed name once with a stable tool-call index/ID, then emit only new argument-text fragments. Retain escape/quote state so boundaries inside JSON strings are handled correctly. Test malformed and truncated calls and ensure fragments reconstruct the original valid arguments.
Use a linked symbol list and heap ordered by merge rank and position; attach versions to discard stale candidate pairs. Only update the neighbours of a merge. Confirm byte-for-byte token ID equality on long repetitive and Unicode inputs before measuring the 50,000-character case.
Build a DAG over input positions with vocabulary-piece edges weighted by log probability, maximize cumulative score with Viterbi and backtrack the pieces. Match normalization, unknown-token behavior and byte fallback as well as segmentation. Compare with deterministic tokenizers encoding, with any segmentation sampling disabled.

36. The server

Check your understanding

Frontend tokenization, JSON and detokenization no longer contend for the engine’s Python interpreter. The core can keep scheduling/launching while the API processes unpredictable text work.
Invalid request parameters fail one request before mutation. A model-step/core crash may leave shared state inconsistent, so the core fails all affected requests and stops accepting work.
Stops are text patterns crossing token boundaries, requiring incremental detokenization. The core owns IDs and numerical state; the frontend owns decoded text and visibility.
Disconnect cancels response consumption; merging cancels producers; generation’s finally sends abort; the core retires the request and frees KV/adapter references. The frontend removes its request queue.
Once HTTP 200 and response headers are sent, it cannot replace them with 503. Admission must precede streaming; later errors are stream events.
The checkpoint’s authors chose defaults for that model’s behavior. Applying unrelated API defaults can change quality and reasoning behavior; explicit client values still override them.

Stretch exercises

Parse the timeout at admission and set a monotonic deadline covering queueing and generation. On expiry, abort core work and finalize collected text with the exercise’s finish_reason: "length"; if streaming has begun, send the final event rather than attempting a new HTTP response. Test expiry while queued, during decode and on disconnect.
Refill each key’s token bucket over time and reserve prompt_tokens + max_tokens before admission. Reject insufficient reservations with 429 and a retry delay based on the deficit/refill rate; optionally refund unused output allowance at completion. Make reservations/refunds atomic so concurrent requests cannot overspend.
On SIGTERM mark readiness false, reject new admissions with 503, and keep existing streams/core work alive until they finish. Bound shutdown time and abort/release remaining requests if the bound expires. Test an in-flight response, a new rejected request and final process/resource cleanup.
Preserve request IDs, bounded backpressure, abort/error propagation and readiness handling when replacing queues with sockets and msgpack. Define the transport’s message ordering and payload contract explicitly. Compare frontend/core latency and CPU cost at the same 256-stream load; serialization changes alone do not guarantee lower overhead.

37. Speculation in serving

Check your understanding

Uncomputed-token counts, ragged verification rows, paged block allocation/trimming and per-request sampling settings already express speculative work. Verification changes commit rules rather than building a separate engine.
A deterministic proposal has q(d)=1, so acceptance is p(d). Acceptance contributes mass p(d) to d; rejection samples the normalized remaining target mass, reconstructing p exactly.
If the request had L known tokens before verification, accepting two drafts leaves L+2 computed positions; the new correction/bonus token remains uncomputed. Blocks entirely beyond the committed computed prefix can be freed; its last partial block remains private.
Drafts were chosen for a particular next step/history. Preemption changes when/how that step is computed, so stale proposals must be regenerated.
A same-position draft runner has separate KV storage with compatible logical block addresses, while the target’s block manager owns allocation decisions. Draft state must be synchronized/restored for newly adopted prefixes.
Larger batches already reuse weights well and can be compute-bound. Extra verification rows, rejected work, draft cost and buffers may cost more than the steps they save.

Stretch exercises

For each k, estimate accepted-plus-correction/bonus tokens and divide by measured/predicted draft-and-verification time at the current batch size. Include k=0, smooth per-request acceptance and probe alternatives. High load can make k=0 best even when acceptance stays high.
Bucket by both request count and k, giving (1+k)*bucket_size verification rows with fixed buffers. Mask padded requests and keep allocation, rejection/commit decisions and ragged exceptions outside capture. Test output/cache equivalence and measure whether padding and verification compute outweigh replay savings.
Train the fused embedding/target-hidden input projection and draft layer on aligned target histories, with compatible vocabulary and positional semantics. Use a separate draft KV pool addressed by the target’s logical block tables and restore it after rejection/preemption. Compare held-out acceptance and total latency with Medusa; teacher-forced training quality alone does not measure rollout acceptance.
Flatten the draft tree, assign depth positions, and let each node attend only to the committed prefix and its ancestor chain. For greedy decoding, traverse target-agreeing edges and commit the longest path plus its target correction/bonus, retaining only committed KV. Exact stochastic tree verification needs an explicit probability-correct acceptance rule beyond choosing the longest path.

38. GGUF

Check your understanding

GGUF includes model/tokenizer metadata and per-tensor GGML types/layouts. A standalone safetensors weight file normally relies on separate config/tokenizer files.
Memory mapping avoids eagerly copying a large file and permits tensor views/lazy page access. Some runtime repacking still needs allocated device/host storage.
A byte pairs weights j and j+16 within a 32-weight block, in low/high nibbles. Pairing adjacent weights scrambles each block while retaining plausible shapes.
Compressing sub-block scales reduces metadata overhead while adapting to local ranges. One scale for the whole legacy block cannot fit those ranges as flexibly.
Q4_K is a tensor encoding. Q4_K_M is a model-level recipe mixing encodings for selected tensors, not a new block layout.
They can be repacked into codes plus affine group scales/offsets. That shared layout simplifies dispatch, but expands some packed scale metadata and can use more memory than native Q4_K.
Llama conversion permutes Q/K rows for its interleaved rotary convention. A split-half decoder must invert that permutation; the supported Qwen conversion convention already matches its row pairing.

Stretch exercises

Compare each GGUF variant with BF16 on identical prompts and subsequent teacher-forced tokens, using a stated KL direction such as KL(p_BF16 || p_quant). Confirm tokenizer IDs and architecture conventions first. Q8_0 often has lower divergence than Q4_K_M, but report measured per-token means/tails rather than assuming a fixed quality gap.
Decode each super-block’s scales/mins and packed codes in registers and multiply without repacking into expanded affine metadata. Native Q4_K is 144 bytes per 256 weights, or 4.5 bits/weight; the chapter’s affine representation costs about 6. Validate exact layout reconstruction and compare achieved bandwidth, residency and output error.
Implement the tokenizer’s normalization and space marker, then repeatedly merge the highest-scored eligible adjacent piece with the specified tie rule; apply byte fallback when no vocabulary piece covers input. Scores are not BPE merge-rank entries. Compare with the checkpoint’s tokenizer on spaces, Unicode and added tokens.
Weight reconstruction errors by calibration-derived per-column importance, fit each sub-block’s affine scale/minimum, and account for quantization of the stored scales/minima themselves. Evaluate on held-out activations and teacher-forced KL with the same model recipe. Lower weighted calibration error does not guarantee every prompt improves.

39. Quantized checkpoints

Check your understanding

GPTQ uses activation-derived curvature to estimate how column errors affect outputs. It compensates remaining columns for each committed quantization error, rather than rounding independently.
Act-order changes column processing/group assignments; g_idx identifies the resulting groups. Reordering codes and the corresponding activation columns makes runtime groups contiguous while preserving the linear operation.
Scaling salient columns gives them more effective quantization resolution. Inverse input scaling can be absorbed into the preceding normalization/projection; it isn’t free if left as a separate runtime operation.
Exponent bits have finite range and spacing. Block scales place local values in a useful portion of that range, avoiding overflow or excessive rounding of small values.
NVFP4 uses smaller blocks and more flexible scale representation/global scaling than MXFP4’s power-of-two block scales. E2M1 values alone do not determine error.
Trellis states share constraints/codewords across a sequence, optimizing joint error beyond independent scalar choices. A small fixed-state recurrence and lookup tables keep decoding manageable.
KL compares the full candidate distribution with a reference and is useful when a small/random model’s task loss is uninformative. Perplexity on representative real data remains necessary; neither metric replaces a task evaluation.

Stretch exercises

Use the same 128 calibration sequences and separate evaluation text/task data, keeping group sizes, eligible layers and packing comparable. Record per-layer output error, teacher-forced KL and downstream accuracy. Sensitive layers and the relative RTN/GPTQ/AWQ ranking are empirical; ensure act-order metadata is interpreted correctly before comparing.
Start each tile with the correct trellis state or reconstruct the required prefix state, then replay code transitions and lookups without fully materializing weights. Move the orthogonal transform to activations using the exact normalization, signs and transpose convention. Validate dense equivalence; tile boundaries cannot reset the recurrence arbitrarily.
Factor the activation curvature into blocks and apply error feedback between blocks while jointly selecting trellis codes within each block. Include the trellis state constraints and transform conventions in reconstruction. Compare held-out quality with plain trellis at the same 2-bit budget, including metadata; the extra encoding complexity must produce a measured benefit.
Read the expert blocks and exponent scales with the checkpoint’s interleaved layout and compare sampled dequantized entries to the BF16 conversion. Cover multiple experts, blocks and scale values; check shape/layout before blaming arithmetic. Match BF16 rounding in the reference comparison.

40. More hardware

Check your understanding

A small-batch mat-vec reuses each weight little, so moving weights dominates. Quantization reduces those bytes if the kernel consumes codes directly.
It reads codes, writes an expanded matrix and reads that matrix again, adding traffic and allocations. Dequantize tiles in registers or consume integer codes inside the dot product.
Int8 activations enable fast integer dot-product instructions; local scales restore magnitude. Quantization/scale work and additional rounding are the costs.
Runtime dispatch can choose supported instructions while leaving the rest of the binary portable. Globally compiling for one machine can emit illegal instructions even before dispatch on another CPU.
Several requests touch many distinct experts, often more than eight per layer. The next step evicts experts the following request needs, causing repeated transfers.
Streaming helps when weights don’t fit, compute can hide transfers, or prefill has enough rows to amortize them. Batch-1 decode over a slow host link is usually dominated by moving every layer’s weights again.

Stretch exercises

Keep quantized codes/scales resident, use simd::gemv_q4 directly and reuse a persistent worker pool across the entire Rust decode loop. Verify logits against the Python reference before comparing warmed tokens/s. Match model file, threads, context and sampling with llama.cpp so conversion or workload differences do not masquerade as loop-overhead gains.
Quantize activation tiles to int8 with scales, reuse weight tiles across 64 prompt rows and accumulate integer products safely before rescaling. Use runtime ISA dispatch for VNNI/AMX and handle matrix tails. Compare prompt throughput and output error; tile packing and activation quantization must be included in end-to-end timing.
Inspect the exact llama.cpp revision’s Q4_K layout handling, scale/minimum unpacking, cooperative loads and SIMD-group reductions, then port the relevant structure. Its native packed-format path avoids the simple kernel’s generic unpacking/repacking costs. Validate on actual Apple hardware; performance cannot be established by non-Metal tests.
Predict next-layer expert needs on another stream, copy candidates early and fall back for actual routing misses. Current-layer output is only an approximation to the next router’s true input unless the architecture explicitly supports pre-gating. Measure accuracy of predictions, wasted bytes and exposed copy stalls while preserving exact final routes.

41. Multi-GPU

Check your understanding

Independent output rows need only the common input. Output/down projections sum partial products over split input features, so their contributions must be reduced to reconstruct the full result.
Once KV heads are fewer than ranks, heads are replicated across ranks. Further TP splits query work, not the size of each replicated KV head.
Compare payload/link bandwidth with collective startup/round-trip latency. At roughly 1 MB, fast NVLink often makes startup a large fraction; PCIe transfer can be comparable or dominant. Topology and collective algorithm decide the crossover.
A single batch traverses stages serially and pays transfers. Multiple microbatches/requests in flight fill the pipeline; this increases throughput, not automatically one request’s latency.
When routing is sufficiently sparse/balanced and dispatch payloads are smaller than repeated projection reductions. Expert placement, duplicate top-k destinations and link topology matter.
Cache affinity alone can overload the replica holding a popular prefix. A load term trades some reuse for bounded queues and available capacity.
Export/import moves logical KV via the existing swap-copy contract, alongside tokens and progress. Seeded sampling must carry generator state, not just the original seed, to continue its random stream exactly.

Stretch exercises

Determine each rank’s projection slices from the architecture before allocation, then read only those safetensors rows/columns and shared small tensors. Respect GQA head replication and shard before fusion. Measure peak host and device memory per rank, including transient slices, rather than only final shard size.
Maintain pp disjoint request batches with stage-local buffers and placeholder/progress state, and schedule their stages concurrently. Preserve per-request order and stop/abort cleanup for in-flight work. Throughput improves by filling bubbles, while single-request latency still includes every stage and transfer.
TP splits each layer’s work but pays frequent collectives; PP uses fewer boundary transfers but needs concurrent batches to fill stages. Over PCIe, PP can benefit throughput at sufficient concurrency while TP may reduce compute time enough to help latency. Measure both with the same model, memory constraints and offered load; topology determines the crossover.
Shard large draft-model projections and KV according to their own head geometry if the draft runs across TP ranks. Medusa’s small heads can be replicated or evaluated on rank 0 if the required target hidden states are available there. Broadcast draft IDs/commit decisions and synchronize draft progress; account for hidden-state gathers and vocabulary-head sharding.
Record an event after each layer produces stable KV, transfer on a separate stream and let the decode receiver wait on its layer’s completion. Preserve block ownership and metadata until copies finish. Overlap is bounded by remaining prefill compute and transfer bandwidth; report exposed tail-transfer time and total TTFT against a non-overlapped baseline.

42. More models

Check your understanding

FlatModel checks for Q/K normalization by attribute presence. Defining only a Q identity would select the Q/K-normalization path and require a K module too; leaving both absent expresses the model’s actual contract.
Short-wavelength, fast-rotating pairs already experienced many full rotations during training. They keep local positional resolution; slow pairs are interpolated and the middle region blended.
Its frequencies can change with sequence length. Cached keys rotated under earlier frequencies then disagree with recomputed keys at the longer length, violating prefix/state equivalence.
Each head’s nonrotary key/value is an up-projection of the shared latent. Only the latent and shared rotary key are cached; distinct projections give distinct head scores/outputs.
Rotation depends on token position and does not commute with a fixed up-projection. The rotary query/key term must remain explicit.
Long prefill can benefit from decompressing once and attending with narrower ordinary heads. Decode avoids repeatedly decompressing the whole history by absorbing projections around latent attention.
It bounds the number of expert groups a token can visit. When groups map to nodes, that bounds communication destinations and cross-node routing traffic.

Stretch exercises

Implement every listed Gemma-specific norm, embedding scale, GeGLU and layer-type rule from the matching config, preserving operation order and window masks. Compare embeddings, intermediate residuals and fixed-input logits with the official model. Generic RMSNorm or SwiGLU substitutions can have plausible shapes but wrong values.
Keep separate layer-group pools/tables and reclaim a window block only when no live query or in-flight chunk can read it. For purely windowed layers, storage scales with the window plus boundary/chunk allowance instead of full context; the ideal capacity gain is roughly 32768/window for those layers alone. Full-attention layers and shared overhead reduce the whole-engine gain.
Decompress latents into per-head K/V once for the prompt, apply the rotary term correctly and run ordinary causal attention. Compare with the absorbed latent path on identical prompts and include decompression/scratch costs. The crossover depends on head widths, prompt length and kernels; it is not a universal token count.
Recognize FP8 tensors and their block scale metadata, construct FP8BlockLinear with the exact scale direction and preserve the architecture registry’s names/layouts. Validate selected dequantized blocks and layer outputs before full-model logits. Unsupported shapes or devices need an explicit reference fallback.
Tile many query heads against each shared 576-wide latent row, reuse loaded cache data and accumulate online softmax/output in FP32. Include the distinct rotary and absorbed value projections correctly. Compare against the generic one-KV-head reference across context lengths, head-group sizes and split counts; larger groups can increase register pressure.

43. Beyond text

Check your understanding

Expanded patch rows, not the original marker count, consume decoder/KV positions. Length checks before expansion can underreserve memory and exceed the model limit.
KV addresses and causal order follow flattened rows; M-RoPE follows spatial/time coordinates and resumes text after the maximum coordinate. Its continuation delta connects the two.
The actual projected inputs and position recipe, including encoder/processor effects, axes, sections and continuation. The implementation hashes those values and snapshots them.
A pooling rule selects/averages vectors; retrieval geometry must be learned with the model’s training objective and text conventions.
Padding is not input content. Including it changes the direction and magnitude with batch padding length, making the same input depend on its neighbours.
The base weight projection runs once over the whole token matrix. Shrink/expand corrections select A/B weights per adapter and scatter back to the original rows.
Slots are reused and overwritten. A stable digest of effective weights prevents old KV from being reused after a different adapter occupies the same slot.
It will resume using cached/recomputed activations for the same adapter. Replacing those weights during preemption would change the request’s model mid-history.

Stretch exercises

Match the trained processor’s resize/crop/normalization, patch ordering and merger weights. Compare projected image rows, all rotary axes/continuation positions and fixed-input decoder logits with the official model. Textual answer plausibility alone cannot detect a wrong spatial layout.
Key encoder outputs by image bytes, preprocessing settings and encoder revision; include merger identity if caching projected rather than raw encoder features. Use a byte-bounded LRU and preserve the position metadata needed by the decoder. Measure encoder-cache and decoder-prefix-cache hits separately, including host/device transfer costs.
Group token rows by adapter, reuse the row permutation across layers and compute shrink/expand with tiled matmuls. Restore row order before operations that depend on sequence layout. Compare with SGMV over rank, rows per adapter and adapter count, since grouping overhead can dominate small groups.
Queue embedding inputs up to a token/item limit or maximum wait, then run a padded/ragged batch with correct masks and per-input pooling. Bound queue size and propagate cancellation/errors. Measure throughput and p99 from admission, including the batching wait, at matched offered load.
Capture fixed-capacity adapter buffers and rank buckets, zero unused rank lanes and update weights in place only after every old lease and replay user retires. Keep effective adapter identity in prefix keys. Test slot replacement, mixed/no-adapter rows and rank changes against eager execution to catch stale graph pointers or weights.

44. Benchmarking

Check your understanding

Slower responses delay the client’s next submissions, lowering offered load exactly when the engine falls behind. Open-loop arrivals reveal queue growth.
Dispatch time moves later. Scheduled arrival stays fixed so waiting for the client cap remains visible in TTFT/E2E.
There is one observed arrival timestamp; all five token times inside the chunk are unknown. Report chunk gaps and average TPOT when its prerequisites hold.
There are no post-first-token intervals and its denominator is zero. Missing data is clearer than claiming zero decode latency.
Failure is part of service behavior. Excluding failed requests makes overload appear to meet its SLOs by simply refusing slow work.
Once greedy outputs differ, histories differ too. Teacher forcing holds history fixed so logit/loss/KL differences measure the computation rather than cascading prompt changes.
Quantization changes memory traffic, kernel dispatch and numerical quality. Report it as a separate speed-quality-memory experiment with the same source checkpoint.
Not ready (503), even if health remains good: it is intentionally refusing new work while admitted requests drain.

Stretch exercises

Run three open-loop trials at each rate and plot p99 TTFT, SLO goodput and queue length over time. The knee is where sustained offered load exceeds service capacity and queues begin growing; short bursts alone do not establish it. Keep failures in the population and record trial duration/warmup.
Separate genuinely cold unique prefixes from warmed repeated prefixes, using equal total memory budgets and the same prompt/output lengths. Report cache hit tokens and compute time alongside TTFT/throughput. Cache-disabled or matched-hit controls help attribute gains to reuse versus kernels.
Send equivalent rendered prompts, tokenizer settings, image preprocessing and generation budgets to every engine. Report vision-encoder time, decoder prefill and end-to-end TTFT separately. Different templates or image resolutions change the workload and invalidate a direct speed comparison.
Kill the core and disconnect the proxy during controlled bursts, then measure time until every affected client fails or completes. Verify readiness drops, admission stops appropriately and KV/adapter/queue resources are released. Bound client timeouts so hung streams cannot silently disappear from the results.
Use one source checkpoint and identical teacher-forced histories/task protocol, then report quality, peak/resident memory, latency and throughput for INT4/FP8. Plot a trade-off curve with failed requests and SLOs included. A faster variant with worse task performance is a different operating point, not an unconditional winner.