# opus prefill-opt: K=512 IQ2_XS prompt_chunk on the 24 GB M5 Air

author: `prefill-opt-opus` (claude-opus-5-5). patch: `/tmp/opus_prefill_opt/prefill_opt.patch` (applies cleanly with
`git apply` on the current `vendor/strata` working tree, verified on a scratch copy `/tmp/opus_prefill_opt/applytest`,
built in `/tmp/opus_prefill_opt/build_final`). nothing tracked in `vendor/strata` was edited.

## tl;dr

| prompt | metric | base | patched | delta |
|---|---|---|---|---|
| 1024 tok (1 chunk, 1023 rows) | chunk ms (4 interleaved pairs) | 4016 / 4029 / 4066 / 4521 | 3440 / 3459 / 3487 / 3497 | **-14..-24 %** |
| 1024 tok | prefill_ms e2e | 4298 / 4309 / 4358 / 4810 | 3586 / 3604 / 3633 / 3645 | |
| 1024 tok | prefill tok/s e2e | 213 - 238 | **281 - 286** | **+20 %** (median) |
| 244 tok | chunk ms (4 pairs) | 2504 / 2528 / 2634 / 2638 | 2207 / 2320 / 2328 / 2338 | -12 % |
| 244 tok | prefill tok/s e2e | 85.0 - 89.5 | **97.6 - 103.2** | **+13 %** |
| 2048 tok, chunk 1024 | prefill_ms (cool run) | 8196 | 7309 | -11 % |
| 2048 tok, `STRATA_METAL_PROMPT_CHUNK=2048` (env only) | prefill_ms (cool run) | 8196 | **5980 (342 tok/s)** | **-27 %** |
| final shipped patch, re-check (1 pair, warm) | 1024 chunk / e2e ms | 4210 / 4565 | 3699 / 3840 | -12 % / -16 % |
| decode after a 1024 prefill, 64 tok | tok/s (2 pairs, cooled) | 22.06 / 22.52 | 22.73 / 21.45 | noise (+-0.6), no regression |

prefill output is **bit-identical**: `STRATA_METAL_DUMP_MOE` dumps of every layer's `.mixed .ids .wts .yr .hh .sgv` plus
`.R_post47 .inj47` are byte-equal for base vs patched, patched vs the parent's `build-metal` binary, run vs run, and
(bonus) chunk 1024 vs chunk 2048 on the 2048-token prompt. `route_parity`: PASS (0 mismatches both phases).

all numbers: `strata_window_cli`, the parent's `env_k512` (ARENA_GB 6.5, wired, MOCE_BONUS 0.85, IO_THREADS 8, ...),
`STRATA_METAL_PROMPT_MOCE` unset (= 0, exact routing), base = same binary with `STRATA_METAL_PROMPT_PRED=0
STRATA_METAL_PLE_ASYNC=0` (exactly the old code path; its dumps are byte-equal to the parent's binary). runs gated on no
other `strata_window_cli` running; base/patched interleaved because the fanless Air throttles hard (see caveats).

## 1. the diagnosis: cold prefill is an SSD problem, not a kernel problem

- per expert (gate+up+down, contiguous copy) = 36.23 GB / (48 x 512) = **1.474 MB**.
- 1024-token chunk: 410 distinct experts / layer, 14.8 % in the arena -> ~16.8 K experts must come off the SSD =
  **24.7 GB per chunk** no matter what.
- raw SSD: `/tmp/opus_prefill_opt/ssdbw` (random 1.47 MB `F_NOCACHE` preads) = **7.97-8.05 GB/s** at 4/8/16 threads.
- so the floor for a cold 1024 chunk is ~**3.09 s** of pure NVMe time; the GPU work alone is ~2.0-2.4 s (KPROF).
  base sat at 4.0-4.5 s because the SSD was idle part of the time and read experts nobody used.

where the SSD went idle / wasted in base (`stream.cpp at_layer`, prompt branch):
1. `layer_rank_[l+1]` is the static profile: only 50.9 % of lookups were streamed ahead, **34.3 % were synchronous
   misses** (GPU stalled, 1.2-1.5 s exposed wait), and ~1.2 K of the 11.2 K streamed-in experts were never used.
2. the next layer's read-ahead was queued only **after** the wait for this layer's misses returned, so the queue ran dry at
   every handshake.
3. `max_pre` capped the read-ahead at 256 (128 for layers >= 40) of the ~350 non-resident experts a layer needs.
4. the PLE row gather (`gather_batch`, ~150 ms of SSD reads per 1K chunk) ran serially on the host before the GPU got its
   first dispatch.

## 2. what the patch does

### A. exact-ish lookahead of layer l+1's experts (the big one)
in `prompt_chunk`, right after layer l's `route_topk`: `dq(blk.{l+1}.ffn_gate_inp) -> p.Wr2`, `gemm(p.mixh, p.Wr2,
p.logits2)`, `route_topk(logits2 -> ids2)` with MoCE forced to 0. these dispatches sit **before** the layer's cut, so at
`at_layer(l)` the host has T x 10 predicted ids for layer l+1 (same trick decode already uses via `predL_`, but for a
whole chunk and computed on the GPU in ~10 ms total for 47 layers).

`ExpertStream::at_layer` (prompt branch): after pinning/queuing layer l's misses (urgent queue) and **before** waiting,
it counts the predicted ids per expert, orders them most-tokens-first, and submits every non-resident one as a
read-ahead of layer l+1 (no 256 cap). urgent misses still go first (io threads drain `urgent_q_` before `pre_q_`), and
the existing drop logic cancels any unstarted read-ahead layer l+1 turns out not to use. with a prediction set, the old
static-rank queue is skipped. layer 0 (no previous layer) gets the whole profile rank at `hold_region` instead of 256.

effect on the 1024 chunk: streamed ahead 50.9 % -> **76.4 %**, misses 34.3 % -> **9.6 %**, exposed wait
1.2-1.5 s -> **0.56-0.61 s**; experts read 18.0 K -> 18.4 K (the ~1.4 K wrong guesses are the only extra bytes).
total SSD traffic is now ~27 GB in ~3.4 s = ~7.9 GB/s, i.e. the drive is saturated.

knobs: `STRATA_METAL_PROMPT_PRED=0` = old behaviour. like the static read-ahead it stands down for a MoCE prompt
(0 < `STRATA_METAL_PROMPT_MOCE` < 100), whose routing reads residency. tested and rejected `STRATA_METAL_PROMPT_PRED_MIN=2` (only experts
predicted by >= 2 tokens): fewer wasted reads (14.3 K streamed) but misses 17 %, chunk +40..110 ms slower; misses cost
more than waste because they idle the GPU. (the PRED_MIN knob is in the patch, default 1.)

### B. PLE gather beside the GPU
`gather_batch` runs on a `std::thread` started before the graph is submitted; `Window::at_cut(0)` joins it (the GPU is
parked at layer 0's handshake anyway), and the PLE projections (`act16(emb)`, `Wpk`/`Wpv` gemms) are emitted after layer
0's `moe_out`, which is still before their only consumer (the l == 0 PLE block). ~80 ms off e2e prefill (3586-3645 vs
3622-3658 ms, sync vs async on the same build). `STRATA_METAL_PLE_ASYNC=0` = old order; non-streaming path unchanged.
the trace line now splits host time: `prompt chunk ... (embd 3 ms, ple 0 ms, graph build 2 ms)`.

## 3. bit-exactness / KLD

- prefill state: byte-equal dumps as listed in the tl;dr (all 48 layers, all rows). the lookahead only changes *which
  slot an expert sits in and when it is read*; routing in the prompt path never reads residency when PROMPT_MOCE = 0.
- teacher-forced KLD (`--logits-out/--logits-ref`, prompt_2048 file, 1024 prompt + 64 tf tokens):
  - with decode MoCE off (`MOCE_BONUS=0`): base vs **base** (identical config, run twice) mean KLD 0.0012, top-1 95.3 %;
    patched vs base 0.0011, top-1 96.9 %; the parent's binary vs base 0.0079 (single run). i.e. **decode itself is not
    run-to-run deterministic** even with MoCE off (residency/timing-dependent paths: EXPERT_SPLIT / SPLIT_TURN read-
    ahead timing). that is pre-existing and independent of this patch, but worth knowing before anyone calls a
    decode-side KLD "bit-exact".
  - with decode MoCE 0.85 vs the MoCE-off reference: base 0.074 / 0.067, patched 0.062 / 0.064 mean KLD (noise-level;
    if anything the patched arena state is a hair better because the streamed-in experts are the ones actually used).
- `route_parity` (built from the patched tree): K=512 mismatches 0, K=256 + MoCE 0.85 mismatches 0, PASS.

## 4. avenues that did not pay (measured, `STRATA_METAL_STREAM_FAKE=1 STRATA_METAL_PKPROF=1`)

KPROF of the patched 1024 chunk, FAKE stream (cool-ish machine): GPU 1967 ms over 2492 dispatches:
`gemm_f16_128_64_4` x517 534 ms, `mm_gu3_22 (IQ2_XS)` x34 442 ms, `mm_down3_42 (Q2_0)` x48 165 ms, `gdn_scanr_4` 138 ms,
`mm_gu3_16 (IQ3_S)` 130 ms, `qsa_flash` 113 ms, `gr_norm_h` 106 ms, `gemm_hc_up` 73 ms, `mm_combine` 42 ms.
(the extra 47 router gemms + 47 route_topk of the lookahead cost ~10 ms of that.)

- **B (expert tiles)**, probe builds with the plan's TM switched:
  - TM 32 + `mm_gu4_*_2_32_32_4` + `mm_down3_*_32_64_32_2`: gu 447 ms (vs 442), down 212 ms (vs 165) -> worse.
  - TM 32 + `mm_gu3_*_32_64_32_2` + down 32_64_32_2: gu 804 ms, down 235 ms -> much worse.
  the IQ2_XS gate+up is at ~5 TFLOPS effective, bound by the codebook decode, and the 16-row tile keeps more threadgroups
  in flight; bigger M tiles lose occupancy faster than they save decode. keep `16_32_32_4` / `16_64_32_2`.
- **C (dense)**: the dense/shexp/router GEMMs are ~3.4 GMAC/token -> 6.9 TFLOP per 1K chunk in 534 ms = **~13 TFLOPS fp16**
  on the 10-core M5; there is no 30 % left there. `dq_f16` is < 50 ms total. not worth a kernel rewrite.
- **why GPU wins barely matter cold**: after A the chunk is ~3.45 s vs an SSD floor of ~3.1-3.4 s (27 GB at 8 GB/s,
  incl. wasted guesses) and GPU ~2.0-2.4 s; the remaining 0.56 s exposed wait is the drive. GPU work only becomes the
  bottleneck when the Air is thermally throttled (see caveats) or on warm/repeated prompts.

## 5. the next lever: chunk 2048 (env, not in the patch defaults)

bytes per token is the only thing left, and chunk size is how you cut it: a 2048 chunk touches ~470 distinct experts a
layer once instead of ~410 twice. on the 2048-token prompt: experts read 32.6 K -> **18.5 K (-43 %)**, prefill 7309 ->
**5980 ms (342 tok/s)** cool (7747 vs 8506 hot). **bit-identical to chunk 1024** (dumps byte-equal). cost: transient prompt
scratch 1.36 -> 2.29 GB (virtual prefill peak 9.7 -> 10.6 GB against an 18.8 GB cap; scratch is unmapped after
prefill, phys_footprint identical afterwards).

recommendation: set `STRATA_METAL_PROMPT_CHUNK=2048` in the K512 serve env (`scripts/serve_k512_local.sh`). i did **not**
flip the default (`stream_chunk = phys >= 96 GB ? 4096 : 1024`) because with an *automatic* arena size (`am.arena_gb < 0`)
`prompt_gemm_need_bytes()` is subtracted from the arena budget, so a bigger default chunk would silently shrink the decode
arena on 24 GB Macs; with `STRATA_METAL_ARENA_GB=6.5` pinned (as in the K512 config) it costs decode nothing.

## 6. caveats

- thermal: the fanless Air throttles after a few minutes of back-to-back prefills; the same patched 1024 chunk measured
  3.44 s cool and up to 3.8-4.3 s hot, decode 22.7 cool vs 16-17 tok/s hot. all A/B numbers above are interleaved
  pairs; absolute numbers depend on how warm the machine was. the KPROF tables are from the middle of a session (GPU
  1.97 s here vs the 2.3-2.4 s the parent recorded).
- the parent's numbers (3.8-4.1 s base chunk) and mine (4.0-4.5 s base) agree within that thermal spread.
- a concurrent `build-metal/strata_window_cli` run (long-ctx sweep) contaminated two early measurements; those are not
  used anywhere above.

## 7. files

- `/tmp/opus_prefill_opt/prefill_opt.patch` - the drop-in (window.cpp, stream.cpp, stream.h; ~230 lines of diff)
- `/tmp/opus_prefill_opt/ssdbw.cpp` - the NVMe bandwidth probe
- `/tmp/opus_prefill_opt/kld/` - teacher-forced logits + prefill dumps used for the exactness checks
- `/tmp/opus_prefill_opt/strata` - the working copy with the probe hooks (`STRATA_PG_TM`, `STRATA_PG_GU_KERNEL`, ... were
  removed again before the patch was cut)