7.7 kB 103 lines
2267f2a8
Round 28 blitz: mtp.fc_hidden [t][c] stream batching (-52.1% n=2 / -29.3% n=3), ple.metal tg_sum barrier halving (-1.5..-5.8%), sample.metal tg_argmax unroll + router_route_8
dawn
⋯
audits/opus_r28_mtp.md
sourcerendered
download raw history
at main Markdown

opus r28 track a: mtp multi-row fc_hidden batching, mtp.metal unrolls, gm_cat audit #

patch: /tmp/opus_r28_mtp/opus_r28_mtp.patch (git -C vendor/strata apply --check passes against the current dirty tree). files: metal/src/engine/window.cpp (mtp_encode, the fc_hidden block + CombNormArgs only), metal/src/engine/msl/mtp.metal, metal/tests/comb_fold_parity.cpp. mtp_chain.metal is untouched: every fixed-trip loop in it already has unroll(full) / unroll_count(4). scratch tree + build: /tmp/opus_r28_mtp/{strata,build}. all numbers are from Apple M5 (the 24 GB air) with the live :8080 server up, so expect some noise.

1. fc_hidden batching (the change) #

before: for n > 1 (catch-up of a accepted drafts, and mtp_prefill chunks), mtp_encode ran one fc_hidden GEMV per stream c (HC = 4 of them), each with n rows, plus 3 mtp_xprep dispatches. so the 6.95 MB Q8_0 weight got read 4 times whatever n was. h was laid out [c][t], with h_stride = MAXT*N.

after: the n*HC streams become GEMV rows in [t][c] order, up to MAXT = 8 per GEMV.

  • mtp_rms_xprep always splits now. it prepares the first plane_tok (= MAXT) stream rows and skips the rest (new guard t >= plane_tok), and it keeps the normalized rows in mhn_ only when n*HC > MAXT.
  • after that, each further 8-row slice is one mtp_xprep (stride N, from mhn_) plus one GEMV.
  • mtp_combine_gnorm reads h + (t*hc + c)*n. h, R and xn now share one layout, and h_stride is gone from CombNormArgs (on both host and kernel).
n GEMVs before → after weight reads dispatches (rms+xp+gemv+comb) cold GPU time, hidden branch
1 (draft step) 1 (4 rows) → 1 (4 rows) 1 → 1 3 → 3 70.4 → 70.1 µs (same)
2 (a=2) 4 × 2 rows → 1 × 8 rows 4 → 1 9 → 3 211.6 → 101.4 µs (−52%)
3 (a=3) 4 × 3 rows → 8 + 4 rows 4 → 2 9 → 5 220.1 → 155.7 µs (−29%)
4 4 × 4 → 2 × 8 4 → 2 9 → 5 228.9 → 190.0 µs (−17%)
8 (prefill chunk) 4 × 8 → 4 × 8 4 → 4 9 → 9 365.1 → 366.1 µs (same)

comb_fold_parity 0 mtp bench measures this: the weights cycle through 16 copies so they come from DRAM, the dispatches run in a barrier chain as in the window, and each number is the median of 9 alternated pairs. two runs agreed to within 1%. you can't get n = 8 below 4 GEMVs without a GEMV that takes more than 8 rows: gemv.metal's NTOK goes up to 8. so the "zero xprep" n = 8 idea from the brief isn't possible at MAXT = 8. its dispatch count stays the same (3 xp + 4 gemv as before).

tiles (gemv_parity --tiles, Q8_0 2560→2560, cold) #

rows 32,2,4,1 16,1,4,1 16,2,4,1 32,1,4,1 8,1,4,1
4 59.9 57.2 62.0 70.7 59.0
8 113.9 86.0 109.7 124.8 86.5
3 56.0 / 66.0 54.7 / 58.7 58.7 / 71.0
2 54.6 / 56.9 54.8 / 56.1 55.5 / 57.2

the existing tile() choice (rows >= 4 → 16,1,4,1) is already the fastest at 4 and 8 rows, so tile() is unchanged.

bitwise (comb_fold_parity 0 mtp [baseline], ALL PASS) #

  • in every case n = 1..8, the batched R and xn are bitwise equal to n separate n = 1 chains. every fc_hidden GEMV now has ≥ 4 rows, which puts it on one tile (16,1,4,1, never gemv_direct_x), so a stream's h no longer depends on n.
  • at n = 1 and n = 4..8, the output is bitwise equal to the old per-stream chain, because it's the same tile.
  • at n = 2 and 3 the output differs from the old chain by max |Δxn| 9.5e-7. the old chain ran 2–3 rows on 32,2,4,1, and now these n match the n = 1 draft path. that's the one intended numeric change, and it only affects the catch-up K/V of the accepted cells.
  • with the original mtp.metal as baseline (/tmp/opus_r28_mtp/orig/metal/src/engine/msl/mtp.metal), the per-stream chain on the old kernels with the old [c][t] layout is bitwise equal to the same chain on the new kernels for every n. so the kernel edits (the split guard, the combine index, the unrolls) don't change any value.

not run: a full-model check, which the brief forbade (:8080 is live). strata_metal_kernels and strata_window_cli compile and link in the scratch build with no new warnings. the expected end-to-end effect is about 110 µs per verify round with a = 2 and about 64 µs with a = 3, so a few tenths of a percent of tok/s at ~47 ms/token. someone needs to confirm it in a window_cli A/B once :8080 is free.

2. mtp.metal unrolls #

added #pragma clang loop unroll(full) to the 8-way part[] sums in mtp_rms_xprep and mtp_combine_gnorm. the summation order is the same, and the baseline A/B above confirms it's bitwise neutral. this is hygiene only: no measurable time change, since clang very likely unrolled the 8-trip loop already.

3. gm_cat audit (comb_fold_parity 0 mtpcat, Q8_0 2560→N+N, one token, bitwise PASS on every tile) #

shape tile cold: 2 gemv → cat hot: 2 gemv → cat
k+v 512+512 32,1,4,4 (window) 21.48 → 21.38 (−0.4%) / 22.61 → 22.51 (−0.4%) +0.2% / +0.4%
k+v 512+512 32,1,4,1 +0.9% / +0.8% +3.4% / +1.5%
gate+up 640+640 32,1,4,4 (window) 27.59 → 27.41 (−0.6%) / −0.2% −0.1% / +0.1%
gate+up 640+640 32,1,4,1 +0.6% / +0.8% +0.1% / +0.3%

gemv_parity --cat (K=1) agrees: 512 hot +3.3%, cold −0.0%; 640 hot +0.3%, cold +0.2%.

a correction to the brief: shared gate+up 2560→640 at rows = 1 runs at 32,1,4,4 (gemv_catk), not 32,1,4,1. tile() line 2591 sends every n_out >= 512 to K=4. the r27 "+2.4% hot" only shows up with the K=1 tile, and only hot. at the window's actual K=4 tile, cat is neutral or slightly ahead, cold and hot. verdict: keep gm_cat as it is. (some hot numbers drift 18→64 µs between runs, probably GPU clock state with the live server, so trust the cold columns.)

4. tile() check on the other MTP Q8_0 shapes (rows = 1, cold, µs) #

  • 2560→512: 32,1,4,4 11.67 · 32,1,4,1 11.73 · 16,1,4,1 12.07 · 16,1,8,4 12.27 · 32,2,4,2 12.37
  • 2560→640: 32,1,4,4 15.06 · 32,1,4,1 15.14 · 16,1,4,1 15.19 · 32,2,4,2 15.81 · 16,1,8,4 15.89
  • 2560→2560: 32,1,4,4 54.2 · 16,1,4,1 56.7
  • 2560→12288 (q_proj): 32,1,4,4 261.8 · 32,1,8,2 267.4 · 32,2,4,1 271.1 · 16,1,4,1 272.1 · 32,2,4,2 281.6
  • 6144→2560 (o_proj): 32,1,8,8 120.5 · 32,1,4,4 124.0 · 32,1,8,2 124.9 · 32,2,4,2 126.0 · 16,1,8,4 127.5
  • 640→2560 (shared down): 32,1,4,1 13.9 · 32,2,4,1 14.5 · 16,1,4,1 14.6 · 8,1,4,1 14.9 · 4,1,4,1 15.4

tile() already picks the fastest tile on every shape, so nothing changes. one thing left open: fc_embedding at rows = 3 (a = 3 catch-up) uses 32,2,4,1, and 16,1,4,1 was 1–7 µs faster in two runs. but that GEMV runs in the same wave as fc_hidden, so most of the gain is probably hidden. too small and noisy to change without an in-window A/B.

repro #

cmake -S /tmp/opus_r28_mtp/strata/metal -B /tmp/opus_r28_mtp/build -G Ninja -DCMAKE_BUILD_TYPE=Release \
  -DSTRATA_GGML_DIR=$PWD/vendor/llama.cpp -DMETAL_CPP_DIR=/tmp/opus_r28_mtp/strata/third_party/metal-cpp
ninja -C /tmp/opus_r28_mtp/build comb_fold_parity gemv_parity strata_metal_kernels
comb_fold_parity 0 mtp /tmp/opus_r28_mtp/orig/metal/src/engine/msl/mtp.metal   # parity incl. baseline kernels
comb_fold_parity 0 mtp bench                                                     # cold paired timings
comb_fold_parity 0 mtpcat                                                        # gm_cat audit
GEMV_AB_COLD=1 gemv_parity --tiles Q8_0 2560 2560 8 16,1,4,1 32,2,4,1 16,2,4,1 32,1,4,1 8,1,4,1