11 kB 139 lines
411096fc
Add prefill lookahead router + async PLE (335-346 tok/s @ 2K), 64K/128K decode benchmarks (20.4-20.8 tok/s), and QSA-safe suffix drafting
dawn
⋯
audits/opus_prefill_opt.md
sourcerendered
download raw history
at main Markdown

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)