# 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.