Exercises
Every coding problem in the book, with hidden tests and hints, by chapter and level.
Chapter 4 · Attention from First Principles
- ✓A batched soft lookup with separate keys and valuesWrite soft_lookup_batch(Q, K, V, tau=1.0), which runs the soft lookup for a whole batch of queries at once. Q has shape (m, d_k), one query per row. K has shape (n, d_k) and V has shape (n, d_v). Score with dot… Warm-up
- 4.1dFind the temperature that gives a target entropyYou are given scores s of shape (n,), not all equal, and a target entropy H_target strictly between the entropy of the cold limit and log n. Return the temperature tau for which H(softmax(s/τ)) = H_target, accurate… Research
- 4.1cKernel regression as attention, without NaNImplement nadaraya_watson(xq, xs, ys, bandwidth). It takes queries xq of shape (m,) and training points xs, ys of shape (n,), and returns predictions of shape (m,) from a Gaussian kernel. Write it as a softmax over… Core
Chapter 12 · FlashAttention and Efficient Attention
- 12.1dThe HBM traffic of a full training stepCount what one training step of standard attention moves, for one head of self-attention with d_v = d_k. Implement attention_step_traffic(T, d_k, k=1, b=2). It returns a dictionary with three keys, and every count is… Core
- 12.1bThe head width that would make the scores compute-boundThe scores kernel reads Q and K (2Td_k numbers), writes S (T^2 numbers), and does 2T^2d_k FLOPs, with b bytes per number. How wide would a head have to be for this kernel to reach the ridge? Implement… Core
- 12.1cRoofline times for a whole attention layerCompare standard and fused attention on the roofline. Implement attention_layer_times(T, d_k, heads, peak, bw, b=2, k=1, softmax_flops=5) for one layer of self-attention with d_v = d_k. Here heads counts every… Core
∇
The problem bank is part of Pro
Every problem from the book, with hidden tests and hints.