Jane Street Attention Kernels interview questions

An error that grows with context length is a running sum losing increments, and BF16 stops adding 1 to a total at 256. The answer turns the symptom into a bisection, finds the state held in the wrong dtype, and shows the error curve that proves it.

FlashAttention does the same arithmetic as naive attention and is several times faster. The answer is a three-line recurrence and one IO count, and interviewers want both.

With T new queries over S cached keys, the obvious lower-triangular mask is wrong and still runs without an error. The answer derives the right one from absolute positions and names the library default that gets it wrong.

The one-query recurrence is a warm-up. A real kernel carries a vector of maxima and a vector of sums for a tile of rows, and the rescale becomes a diagonal matrix. The answer writes it, runs it, and shows the edge case that turns it into NaN.

The textbook IO count says FlashAttention re-reads K and V once per query block, which at 8K is only 2x better than naive. The answer explains why the real HBM number is 65x lower, and which cache makes the difference.

Storing the attention matrix for backward costs gigabytes per layer; recomputing it costs extra FLOPs. The answer derives the gradient, shows the one scalar per row that makes recomputation exact, and the identity that removes the last N-squared term.

Batch-1 decode at long context gives the GPU eight thread blocks and 132 SMs. The answer counts the waves, picks the split count from a wave-efficiency rule rather than a guess, and derives the merge that makes the result exact.

The indirection adds 4 bytes per 64 KiB of KV, so the byte count barely moves. The answer shows where the real cost hides: dependent loads, tiles that span several pages, and copy engines that want one contiguous box.

GQA shrinks the cache by the group size, but a kernel that loops over query heads reads it back group-size times anyway. The answer packs the group into the matrix dimension and shows what that does to intensity and tensor-core utilization.

FlashAttention-2 reaches well under half of an H100's BF16 peak. The answer explains why in one ratio, exponentials against matmuls, and maps each FA3 technique to the stall it hides.

Three quantization points sit inside FP8 attention, and they fail for different reasons: outlier channels in keys, a probability matrix that lives near zero, and accumulators. The answer simulates each one and says which fix addresses which.

A 4K window over 32K tokens keeps one eighth of the keys, and the speedup over causal attention is not 8x. The answer counts tiles exactly, classifies which need masking, and explains where the rest of the gap goes.

Fusing RoPE saves a read and a write of Q and K, which is real but small in prefill. The precision bug is not small: one tensor in the wrong dtype scrambles positional information by a few hundred tokens. The answer sizes both.

Packing removes padding but puts different requests' tokens side by side in one tensor, so the kernel's tile boundaries no longer match sequence boundaries. The answer shows the indexing and a bug that a causal test can never catch.

At 1M tokens a 70B GQA model needs 320 GiB of KV per request and spends over 90% of prefill FLOPs in attention. The answer lays out a plan that fits on eight H100s, then shows which single change to the model breaks it and what replaces it.

Review a shape-by-shape timing fixture, catch a ragged-cache correctness failure, and calculate whether selective dispatch pays for itself. Includes a small reference you can run without a GPU.