The fastest kernel just moves the bottleneck.

The five phases

A two-stage retriever never scores every document. Stage 1 cheaply narrows millions of docs to a shortlist of ~4096; stage 2 — the kernels of classes 02–05 — scores that shortlist precisely. Everything in this class lives in stage 1, and it is exactly five phases. Click any one:

The phase names are ours; the profiler prints two of them differently, so once and for all: flood is the phase it labels approx, and prune is the phase it labels sort. Here is the whole pipeline as measured, before this class's night of work touched anything (an earlier round had already quantized the flood — chapter 02 tells that part):

These are per-phase medians, so they don't sum exactly to the total median (2.50 + 1.05 + 0.39 + 2.71 + 0.11 = 6.76, not 6.70) — a small honesty the whole class respects: the median of a sum is not the sum of medians. Two phases carry the pipeline — flood (40%) and cdot (37%) — and the rest of this class is a profile-driven walk down that list, re-sorting it after every win.

What a "flood" is

The flood is PLAID's approximate scoring pass — how stage 1 ranks candidates without ever touching an embedding. The trick is a decision made at index time: every document token was snapped to its nearest centroid, an integer code in 0..K. So for pruning purposes a document is nothing but its list of codes.

That makes MaxSim over centroid-snapped tokens a pure table lookup into the very matrix stage 1 already built. cdot is q · centroids.T — every query token scored against every centroid, shape [nq, K]. A document's approximate score is then:

No embeddings, no decompression, not a single multiply — just row lookups and maxes into a matrix that already exists. At fiqa-52k that is ~3.5M row-lookups + maxes per query. Watch it on a toy: a 4-token query, six centroids, and a document that snapped to codes 1, 4, 4, 5:

Now the licence for everything that follows. The flood may be sloppy, because its output only has to rank candidates for a 4096-deep cut: a document survives as long as it lands anywhere in the top 4096. So the whole cdot matrix can be quantized to u8 before the flood ever reads it — and it costs nothing measurable:

Gate passed. u8-quantizing the entire matrix is nDCG-identical to four decimals — two datasets, both quant schemes, and under the stage-2 LUT overlay. The 4096-deep prune cut absorbs sub-LSB rank perturbations exactly, so a sloppier-but-cheaper flood is a free lunch. That is why cdot's output can be bytes, which is what makes chapter 05's kernel possible.

cdot: free parallelism, bit for bit

2.50 ms → 0.64 ms. cdot is one GEMM — [nq, dim] · [dim, K], the query's tokens against all K centroids — and at baseline it ran single-threaded on an otherwise idle 8-core machine. The fix is the most boring kind of speedup, and that is the lesson.

Why is a parallel GEMM bit-for-bit identical when floating-point addition isn't associative? Because we split outputs, never a reduction. Each output element is one dim-128 dot product, and that whole 128-long summation lives inside a single column block — no block boundary ever falls in the middle of a reduction. The threads compute different elements; not one of them re-orders a single add. Parallelism touched the schedule, never the arithmetic.

probe: selection becomes a stream

1.05 ms → 0.11 ms on the M4; 13× on Neoverse. For each query token, probe picks the top-8 IVF cells to search. The old way built a full K-length index buffer (128 KB of writes at K = 32k) and ran select_nth_unstable with indirect comparisons. We never needed the buffer.

Replace "materialize then quickselect" with a single streaming scan that keeps a running threshold — the current 8th-best. Read the K similarities once, in order, one 64-value chunk at a time. Per chunk: compute the chunk max with a branch-free SIMD reduction (~4 vector instructions), and skip the whole chunk unless its max beats the threshold. The scalar per-value path — the actual top-8 insert — almost never runs. Step through it (top-3 here, for screen; 8-value chunks stand in for the real 64):

The common case is a chunk of 64 values retired by ~4 vector instructions plus one branch — and that branch is predictable, because once the threshold climbs, almost every chunk fails it the same way, run after run. Modern CPUs speculate a well-behaved branch for nearly free; it is the unpredictable branch, or the pointer-chasing compare, that stalls a pipeline. Same top-8 set by value; ties stay arbitrary, exactly as select_nth left them.

flood: registers beat memory

Flood 2.29 ms → 0.53 ms. The flood is chapter 02's row-lookup-and-max, ~3.5M of them per query. Its two umaxes were never the cost. We proved that by disassembly — echoing class 05's standing rule: you check the codegen, you never assume it.

The old loop was doc-token-outer, and it paid twice per token: it round-tripped the accumulator through memory on every token, and its inner loop had a runtime trip count (nq is a runtime value), so LLVM kept the trip checks, prologue, and epilogue that a static count would delete. Both costs are pure overhead around two max instructions.

for tok in doc.codes:        // outer
  c = tok
  for qi in 0..nq:          // runtime nq!
    a = load acc[qi]           // ↩ memory
    a = max(a, cdot_q8[c][qi])
    store acc[qi], a          // ↩ memory
// per token: nq load+max+store, plus a
// dynamic-length loop's prologue/epilogue

// rows padded to a 16-multiple stride:
// pad bytes quantize to 0 = max-neutral,
// and add 0 to the final sum. both free.
for chunk in 0..nq_padded/16:   // outer
  acc16: [u8;16] = 0            // one vreg
  for tok in doc.codes:        // inner
    v = ld1 cdot_q8[tok][chunk] // 16 B
    acc16 = umax acc16, v      // ← the only op
// no acc traffic, no tail; codes re-read
// once per chunk (≤2 passes at nq≤32, L1-hot)

Invert the loops to chunk-outer / tokens-inner and give each 16-wide chunk a [u8;16] accumulator that LLVM promotes to a single vector register. Pad the quantized cdot rows to a 16-multiple stride, and the tail vanishes: the pad bytes quantize to 0, which is the identity for both max and the final sum — two free invariants. Per token now: load a code, load 16 bytes, one umax.16b. That is the whole inner loop, and the disassembly confirms it — exactly one umax.16b, no accumulator spill, no dynamic epilogue. The old per-token load/store + trip machinery was ~2/3 of the flood.

The transpose-and-quantize prep that builds this padded, code-major u8 matrix is its own sub-phase (chapter 06 parallelizes it). It is why chapter 02's "cdot output can be bytes" matters so much: the byte layout is the register-resident flood.

gather: the address space is the sort

0.41 ms → 0.10 ms. Gather walks the probed cells' IVF postings and produces the candidate list — deduplicated and sorted — for the flood. Doc ids are dense integers, and that one fact hands you the whole algorithm.

Allocate a bitmap with one bit per doc. Walk the concatenated postings once, setting bits; duplicates collapse for free because setting a set bit is a no-op — an O(n) dedup with no hash, no comparisons. Then scan the bitmap in order: every set bit you encounter is the next candidate, and because you scan low address to high, they come out already sorted. No sort call at all — the address space did the sorting. At 52k docs the bitmap is 52000 / 8 = 6.5 KB, small enough to stay resident in L1 for the whole walk.

Two more changes rounded out the night, and both are honest about being small:

• Parallel transpose + quantize (tq 0.35 → 0.20 ms). Same argument as cdot: the blocks are disjoint and each writes contiguous output, and because min/max over a block is order-independent, the result is bit-identical. Modest (~1.7×), but free — it is the prep that builds chapter 05's padded matrix.
• Sorted probe cells (neutral on the M4). The selected cells came out of a HashSet in arbitrary order, so postings were read in random cell order. Sorting ~100–200 cell ids costs nothing and makes the postings reads sequential. On the M4 the L2 absorbs the scatter, so it measures as noise — kept anyway for server parts, which is the whole point of chapter 10.

Amdahl: the honest headline

Per-phase ratios ranged from 3.6× to 9.5×. The total was 4.0×. That gap is not disappointing — it is Amdahl's law, and it is the single most useful equation for deciding what to optimize next. Here is the whole night, phase by phase:

The composition is not an average of the ratios — it is:

Plug the shares and per-phase ratios in and you get back ~4.0×. And it exposes the trap in chasing ratios: the probe's 9.5× was the biggest ratio of the night — and worth the least. Probe is only 16% of the pipeline, so even an infinitely fast probe caps the whole pipeline at 1 / (1 − 0.16) = 1.19×. Toggle each real win on and off and watch which slices actually move the number:

Backwards, on purpose

Read the whole optimization campaign end to end — classes 02 through 06 — and it worked through the pipeline in reverse. That was not sloppiness, it was the profile talking. Each win created the measurement that justified the next.

• Stage 2 first (classes 02–05). The fused kernels made rescoring cheap — which made stage 1 the dominant cost. You cannot see stage 1 as the problem until stage 2 stops being one.
• The flood next. Inside stage 1, the flood was 62–67% of the time, so it went first (the earlier q8-flood + mechanics commits). Attacking the biggest slice, exactly as chapter 07 prescribes.
• Then everything it was hiding. Knocking the flood down surfaced cdot, probe, and gather as the new leaders — this night's work. The profile you fix is never the profile you inherit.

There is no way to have started at cdot. At the start of the campaign it was a rounding error next to the float rescore; it only became worth a parallel GEMM once four other things got out of its way. "Optimize back to front" is just "re-profile and follow the biggest number" run to its conclusion.

The kill

The best thing we did to the flood was not write a change to it. One idea looked clever, clean, and losslessly free — and died in ten minutes, before a line of Rust.

The lesson is the ratio of effort to certainty. That kill cost a ten-minute counting script and bought total confidence the idea was dead; writing it would have cost a kernel, a bit-identity test, and a benchmark to reach the same "no". The cheapest optimization is the one you kill before you write it — and the way you kill it cheaply is to measure the thing it depends on first.

Never trust the M4 alone

Every number so far is an Apple M4 — and the M4 is the least favorable platform these changes have. Its huge, low-latency L2 quietly absorbs exactly the scatter that server parts pay full price for, so it systematically understates every scatter-elimination win.

The probe scan is 6.7× on the M4 (scifact; 9.5× at fiqa-52k) but 13× on a Neoverse N2 (2.23 → 0.17 ms). The sorted-probe-cells change looks like dead weight on the M4 — pure noise — yet it is exactly the sequential-postings win that server parts want. Judge either on Apple silicon alone and you ship the wrong default. Here is the whole night confirmed across four CPUs:

Shared CI VMs — read the ratios, not the microseconds (the instances differ between runs, but the ratios sit far above runner noise, and every parity gate was green on all three). The M4's 4.0× is the smallest absolute stage-1 time and among the ratios; the server parts had more scatter to eliminate, so they had more to gain.

Seven questions

Instant feedback, explanations included. Miss one? The chapter number is next to it.