11 kB 139 lines
bec6ab3b
Round 1 blitz: MTP batched fc_hidden + sink=64 + valid-run KV, single-rw mm_gu3, multi-tok M5 gemv/moe tiles, prompt publish_ready (32.7-40.8 tok/s)
dawn
⋯
audits/opus_mtp_turbo.md
sourcerendered
download raw history
at main Markdown

opus-mtp-turbo: Track 1 report (MTP drafter: indexer, sink window, low-bit weights, sampling at temp > 0) #

Patch: /tmp/opus_mtp_turbo/opus_mtp_turbo.patch (617 lines, paths relative to vendor/strata; apply with git -C vendor/strata apply /tmp/opus_mtp_turbo/opus_mtp_turbo.patch; --check passes against the live tree as of this writing, which other sessions are editing). Touches metal/kernels/qsa.metal, metal/src/engine/window.cpp, metal/src/strata_window_cli.cpp. Nothing here ran on the GPU: every speed number is a byte-count floor, not a measurement.

What was verified (CPU only) #

check result
clang++ -fsyntax-only -std=gnu++20 -arch arm64 -DSTRATA_GPU_METAL=1 -Wall on patched window.cpp / strata_window_cli.cpp clean; same 4 pre-existing warnings as the unpatched file
patched qsa.metal compiled through MTLDevice newLibraryWithSource + qsa_prep pipelines with QSA_ROW0 unset and = 24 (/tmp/opus_mtp_turbo/mcheck.mm; no GPU work dispatched) both pipelines build
exactness of the new temp > 0 verification, Monte Carlo, 400k sequences of 3 tokens, toy 6-token model with history-dependent logits, temp 0.6 / top-p 0.95 / min-p 0.05 / rep 1.08 / pres 0.4 (spec_test.cpp) TV(ref, ref2) = 0.0027 (noise floor), TV(ref, exact) = 0.0029, TV(ref, current greedy-match) = 0.0598
same, drafter = the model's own greedy chain (spec_test2.cpp) exact 1.83 tok/round, TV 0.0028; greedy-match 3.00 tok/round, TV 0.79
sink / window cell bookkeeping, randomized sim of streaming prompts + rollbacks + decode rounds, checks that every cell the draft attention reads holds committed K/V new logic: 0 invalid reads in 3 x 36k rounds; original logic: 407k invalid cell reads in 36k rounds
Q8_0 / Q5_0 / Q4_0 (Lloyd) on the real MTP BF16 tensors, 256 rows each, GEMV error on Gaussian x (quant_test.cpp) table below
CPU cost of one sampling-distribution build (copy + top-k + penalties, V = 248320) 0.23 ms per row

Findings, item by item #

1. Dead indexer compute in mtp_encode: real, removed (small win) #

Confirmed: mtp_encode never calls qsa_score / qsa_topk; mqids_ is filled on the CPU. The indexer state that qsa_prep writes for the drafter (qidx_, mtail_[0], mdead_[0], mpooled_[0]) is read by nobody: the only other references are checkpoint/session save/load and SCR relocate, which copy bytes around without using them. Correction to the brief: qsa_prep's grid rows are 0..23 q heads, 24..27 K/V, 28..31 indexer q heads, and x = 32 the indexer-key threadgroup (not "x = 32 is kcur/vcur").

Patch:

  • index_qk_proj GEMV dropped (and not quantized at load).
  • qsa_prep dispatched with grid x = 28 on draft steps (q + K/V), and on catch-up / prefill rows (no queries read) with a new function constant QSA_ROW0 (index 28, unused elsewhere) = 24 and grid x = 4: K/V rows only. The main window's qsa_prep is unchanged (constant undefined -> 0).
  • Saves 1.74 MB (640 x 2560 Q8_0) per encoded row group, 29 of 33 qsa_prep threadgroups per catch-up row, and the serial 1-threadgroup indexer-key loop. That's ~1% of a draft step's bytes, not "an entire step".

2. Sink + window: implemented, and it uncovered a worse pre-existing bug #

Pre-existing bug (fixed): with SSD streaming, prompt_run_rows calls on_chunk (the drafter's prefill) only for the last prompt chunk. If that chunk has fewer than 256 rows, mtp_prefill skips nothing, so mtp_cell0_ is never moved, and the draft window [cell0 - 256, cell0] covers up to 255 cells of the previous chunk that the drafter never encoded: zeros after a fresh start (zero keys still take softmax mass, with zero values), or the previous conversation's K/V after window_reset (reset does not clear mkq_). For a 12000-token prompt with 512-row chunks the last chunk is 224 rows, so 32 garbage cells sit in every early draft. This alone could explain part of an acceptance drop on long Albedo sessions, independent of the sink.

Fix: the drafter now tracks its valid run [mtp_cell0_, mtp_cell1_) plus the sink [0, mtp_sink_) through one helper, mtp_wrote(c0, n) (called by mtp_prefill for the parts it writes and by mtp_draft for the catch-up rows + draft 0, which hold committed tokens). A write at c0 marks everything past c0 stale (rollbacks / checkpoint restores to shorter prefixes stay correct); a gap restarts the run. reset() and session_load() clear it.

Sink: STRATA_METAL_MTP_SINK (default 64, 0 = off, max 1024). mtp_prefill also encodes the skipped rows that fall below the sink cap; the streaming prompt_run_rows path additionally calls on_chunk on the first chunk with just its sink rows (64 rows, cell0 = 0), since its residuals are gone by the last chunk. The draft attention's mqids_ is [0 .. min(sink, first) - 1] ++ [first .. cell0], at most 64 + 257 = 321 cells = 6 chunks of 64 (was 5). qsa_attn handles arbitrary cell ids fine (cid[] gathers by id, RoPE was applied at write time at the absolute position, so relative positions stay right); qsa_merge derives the chunk count from pos0 = count - 1.

The "78% -> 61%" figure in the brief is unverified by me (needs the GPU). Measure with SINK=0 vs 64 to separate the two effects (the stale-cell fix is on in both).

3. Q4_0 / Q5_0 draft weights + tiles: partly a wrong premise #

  • mhead_ is not Q8_0. In K512-iq2_xs-metal.gguf, output.weight is IQ4_XS (type 23), and mhead_ is a row gather in the file's type: 40525 x 1360 B = 55.1 MB, not 110 MB. Q4_0 (4.5 bpw) would be larger than IQ4_XS (4.25 bpw). It already runs through Window::gemv on the M5 32,2,4,1 tile (m5 && IQ4_XS && n_out >= 6144). No head change.
  • A bigger, free byte cut was hiding in fc_hidden: the 2560 x 2560 matrix was read once per hyper-connection stream (4x = 27.9 MB per draft step). For n = 1 (every draft step) the 4 streams are now the 4 rows of one GEMV (xprep with stride N, y_stride = MAXT * N lands each row at its stream's mh2_ slot). Each token's sum in gemv.metal is independent of the token count, so the result should be bitwise the same; worth a parity check. Saves 20.9 MB per draft step (~11%).
  • mq8_ low-bit, opt-in: STRATA_METAL_MTP_Q4=1 / STRATA_METAL_MTP_Q5=1 puts q_proj, o_proj, fc_embedding, fc_hidden and the shared expert in Q4_0 (Lloyd, the HC loader's quantizer) / Q5_0 (ggml rule). K/V, router, shared_expert_gate and the hyper-connections stay Q8_0 (gr_read_m only reads all-Q8_0 or BF16; a Q4 HC matrix there would be read at the wrong offset). Both quantizers are lifted out of Window::init's lambdas into shared functions (no behaviour change for the HC path).
  • Tiles: in mtp_encode's gm, Q4_0/Q5_0 wide matrices (n_out > 640, n_in < 6144) on the M5 layout take 32,2,4,1 (the precedent Window::gemv uses for Q8_0 512-row and IQ4_XS on the M5); unmeasured. Q8_0 defaults are unchanged; STRATA_METAL_MTP_TILE=L,R,S,K overrides the wide class for A/B (must be one of GV_SET's combos).

Quantization error on the real MTP tensors (256 rows each, GEMV output rel-RMSE on Gaussian x):

tensor Q8_0 Q5_0 Q4_0 (Lloyd)
fc_embedding 0.58% 4.4% 8.7%
fc_hidden 1.18% 4.4% 5.8%
shared_expert.down 0.68% 5.6% 10.7%
o_proj 0.62% 4.9% 9.1%
q_proj 0.69% 5.6% 10.3%

Byte floor per full draft step (n = 1), from tensor sizes: before ~186 MB (fc_emb 7.0, fc_hidden 4 x 7.0, 3 HC blocks 21, K/V 2.8, indexer 1.7, q 33.4, o 16.7, router 1.4, shared 5.2, 10 Q2_0 experts 13.8, head 55.1). Patch default: ~163 MB (-12%). MTP_Q5=1: ~139 MB (-25%). MTP_Q4=1: ~131 MB (-30%). At ~120 GB/s that's 1.55 -> 1.36 / 1.16 / 1.09 ms per step: pure bandwidth floors, dispatch overhead not included. Whether Q5/Q4 pay off depends on the acceptance they cost, which only the GPU can tell; try Q5 first.

4. temp > 0 verification: the current scheme is not a sampler; replaced with exact speculative sampling #

Current serve_loop at temp > 0: drafts are kept while they equal output[i], the unpenalized greedy argmax, and only output[keep-1] is sampled. So every accepted draft token is a greedy, penalty-free decision: with a good drafter, most of the text is greedy decoding, and rep/presence penalties apply only to the bonus token. In the toy test with a perfect greedy drafter this gives TV 0.79 from true sampling (3.0 tok/round). The "high acceptance at temp 0.6" is partly this artifact.

Second bug: the bonus token was sampled before cached_ids got the window's tokens, so the penalties missed input[0..keep-1], including the token right before (the one-token path pushes cur_tok first). Fixed on both paths.

New default: spec_sample (Leviathan et al. 2023 sec. 2.3, deterministic drafts so q is one-hot): row i's distribution p_i is exactly sample_logits' (sample_dist factored out: top-k, penalties over history + input[0..i], temp, min-p, top-p); draft x is accepted with probability p_i(x); on rejection, sample p_i with x removed (= norm(max(p - q, 0))); a fully accepted window samples its last row. Tokens are distributed exactly as one-at-a-time sample_logits (TV at noise floor above). sample_logits itself is unchanged in behaviour. STRATA_SPEC_SAMPLE=greedy keeps the old greedy-match rule (with the history fix). The greedy stop-token shortcut is gone on the exact path (EOS is sampled like any token, as the one-token path does). Cost: ~0.23 ms CPU per verified row (<= 4 rows).

Expect tokens/round at temp 0.6 to drop vs. greedy-match (accepted ~= E[p(draft)] instead of P(draft == argmax)); that's the price of the output being the distribution you asked for. Typical-acceptance (Medusa) or margin rules would get the speed back but are lossy again, so I didn't ship one. If the speed hit is too much, the honest options are temp 0 for drafted spans or a lower temperature, not a relaxed verifier.

Suggested GPU checks (for the parent; I did not run any) #

  1. Parity: temp 0, STRATA_METAL_MTP_SINK=0, a prompt shorter than one chunk: draft tokens/acceptance should match the old binary exactly (fc_hidden batching, indexer removal and QSA_ROW0 should not change bits).
  2. Long persona prompt (~12k tokens, streaming), temp 0: acceptance old vs new SINK=0 (stale-cell fix only) vs SINK=64.
  3. MTP ms per round: default vs MTP_Q5=1 vs MTP_Q4=1 (+ acceptance), and STRATA_METAL_MTP_TILE A/B.
  4. temp 0.6: tokens/round and tok/s, default (exact) vs STRATA_SPEC_SAMPLE=greedy.

Not done / left out #

  • Q4/Q5 for the drafter's hyper-connection matrices (21 MB/step): needs gr_read_m to pick gr_down_q4/q5 the way the main window does; the HC mixing is the most precision-sensitive part, so it should be measured separately.
  • Batching fc_hidden for n > 1 (prefill/catch-up): the row layout doesn't match mh2_'s; small win.
  • Nothing measured on the GPU, by instruction.