gem/docs/research/kernel-roadmap.md

688 lines
51 KiB
Markdown

# Gemma 4 M4 Max kernel roadmap
This is a native-engine roadmap. Reference measurements may use the adjacent llama.cpp checkout,
but promoted kernels live in this repository and cannot call ggml at runtime.
## K0: loader and contract
- Strictly parse GGUFv2/v3 with overflow, truncation, duplicate, alignment, type, and extent checks.
- Accept only the exact dense 12B/31B metadata topology and tensor manifests.
- Validate both published Q4_0 headers without downloading their data sections.
- Gate native binaries against accidental llama/ggml linkage.
Exit: synthetic and sanitizer tests pass; authentic 12B/31B headers report 667/833 tensors.
Status: complete. The authentic 31B mapping also binds all 60 executable layer descriptions.
## K1: reference correctness core
- Implement Gemma 4 tokenizer metadata ingestion and golden tokenization/detokenization vectors.
- Implement scalar/F32 RMSNorm, Q4_0 and Q8_0 dequantization, matrix-vector, GELU-gated FFN,
proportional/default RoPE, causal SWA/global attention, residuals, and logit soft-capping.
- Build a one-layer synthetic fixture and compare every named intermediate with a high-precision
reference. Then compare first-token logits from real weights.
The CPU path is a correctness oracle and diagnostic fallback, not the intended 31B serving path.
Status: scalar operators, fixed decode layer, and wrapped SWA KV behavior pass synthetic tests.
Real-weight first-token differentials now cover both official Q4_0 sizes against the adjacent Metal
reference. For the prompt `Hello`, both tokenize as `[2, 9259]`; all 262,144 logits are finite and retain
the same argmax and top-10 set. The 12B/31B correlations are 0.9999958/0.9998622 with NMSE
8.38e-6/1.93e-5. Per-intermediate real-weight dumping remains a diagnostic option rather than a serving
gate because full-vocabulary output already exercises every fixed layer and projection.
## K2: Metal decode path
In order:
1. direct Q4_0 matrix-vector with fused dequantization and F32 accumulation;
2. fused RMSNorm + Q/K projection preparation where traces justify it;
3. Q/K RMSNorm + RoPE;
4. single-token SWA/global attention over the unified-memory KV cache;
5. gated GELU FFN projection pair and residual/norm epilogues;
6. final Q8_0 embedding/output projection and on-GPU top-k reduction.
Status: a full allocation-free synthetic Metal layer passes the scalar C layer oracle, including
Q4_0 projections, Q8_0 tied embedding/logit paths, normalized RoPE, SWA/global GQA attention,
GELU gating, and residual epilogues. The model runner binds one mapped weight buffer and loops the
fixed 48/60-layer schedule. Persistent/no-copy buffers pass range and numerical tests. At the 31B FFN gate shape,
three 200-iteration samples measured 13.35 GB/s with copy-per-dispatch and 107.43 GB/s with
persistent buffers. The +705% diagnostic clears the variance threshold, but is not promoted:
deterministic real-model generation and model-level profiling evidence are still required.
The scheduler can now encode dependent operations into one command buffer with explicit buffer
barriers. At the same 31B FFN shape, persistent span 1 measured 99.65 GB/s and span 7 measured
138.61 GB/s (+39.10%, 2.75% noise floor). Weighted row-wise RMSNorm and NeoX/proportional RoPE
also pass CPU differential tests, including RMSNorm-to-Q4 projection in one command buffer.
The token-spanning scheduler records embedding, all 48/60 layers, and logits into one command
buffer; a synthetic two-layer dependency chain matches the CPU oracle. On authentic 31B Q4_0,
token span produced bit-identical logits to layer span and improved pp from 7.100 to 7.561 t/s
(+6.48%) and tg from 6.946 to 7.427 t/s (+6.93%) over three repetitions. It is now the default;
layer span remains available as the correctness/benchmark baseline.
The Q4_0 decode kernel now computes four output rows per SIMD group, independently adapting the
multi-row geometry from the pinned MIT llama.cpp Metal kernel. The reference MTL0 MUL_MAT suite
passed 1154/1154 cases before adaptation. Native CPU differential tests pass; authentic 31B logits
retain argmax/top-10 with NMSE 2.60e-11 versus the old kernel. Real-shape bandwidth rose from
138.61 to 331.93 GB/s, and model pp/tg rose from 7.56/7.43 to 19.80/18.99 t/s. A two-SIMD-group
threadgroup variant was not promoted: +1.04% tg was inside 3.52% noise and pp regressed 1.67%.
For every kernel record shape, bytes read/written, threadgroup geometry, GPU duration, numerical
error, and tokens/s. Compare against the adjacent llama.cpp Metal backend on the same GGUF and
fixed prompt.
## K3: prefill path
- Batched Q4_0 matrix-matrix kernels sized for 3840/5376 embeddings and 15360/21504 FFNs.
- Flash-style attention specialized to the repeating five-SWA/one-full pattern.
- Avoid materializing masked score matrices and discard prompt rows as soon as their downstream
values are no longer needed.
- Measure 512, 2K, 8K, and hax-shaped repeated prompts independently from decode.
Status: the promoted 64-output x 32-prompt Q4_0 SIMD-matrix operator passes its CPU differential.
Its bank-aware threadgroup swizzle reaches 38,663 rows/s at the 31B FFN shape, +44.78% over the
prior linear layout, and improves the down projection 95.02%. The one-axis model gate improves pp
from 64.17 to 87.17 t/s (+35.84%) while 64-token generation changes +0.23% inside variance. All
262,144 comparison logits are bit-identical. The obsolete linear kernel was removed; batches below
eight still route to the decode matvec kernel and batch 32 remains the default.
Native Metal timestamp profiling shows GPU execution consumes 99.07%/99.33% of wall time for
32/512-token prompts, ruling out CPU command encoding as the current limit. Position-batched RoPE
reduces a 32-token prompt from 11,732 to 8,012 dispatches and a 512-token prompt from 200,312 to
140,792. The authentic gate improves pp from 87.17 to 97.71 t/s (+12.10%) with bit-identical
262,144-entry logits and no generation regression. It is now the production RoPE path.
A non-wrapping position-batched KV-store trial was removed without promotion. Sequential results
looked slower despite fewer dispatches, but a post-reversion control crossed below the candidate,
so thermal/run-order drift invalidated the comparison. Retry only with an interleaved in-process
A/B harness; retain row-wise store-to-attention ordering meanwhile.
The token-parallel 256-thread attention kernel is also promoted. It crosses over at eight live
tokens, improves the 32K full-attention operator 12.85x, and raises authentic 256-token 31B pp
124.78% without a tg regression outside variance. F16 KV is correct and halves cache memory, but
its 256-token model speed stayed inside noise, so it remains a candidate rather than the default.
A tiled GQA-reuse candidate subsequently reached 3.57x over that promoted kernel at the 32K
full-attention operator and 1.99x at the 1K SWA operator, but failed the authentic repeated-layer
gate: error was already material at 1K and logits became entirely NaN by 1,792/2,048 tokens.
A global-softmax split-score alternative erased the speedup. Both were removed. Future flash-style
work must pass wrapped-ring, repeated-layer, and 2K authentic-logit gates before model benchmarking;
isolated operator throughput is not promotion evidence.
GPU-timestamp profiling now separates submission overhead from kernel time. Command batching is a
large 1K lever (79.36% to 98.52% GPU wall fraction from span one to twenty) but is exhausted at 32K
(99.84% GPU wall fraction at span five). A numerically exact shared-weight value tile changed 32K
throughput by only +1.00% and was removed under the 3% gate. The subsequent exact split-GQA path
keeps one global softmax per query head and reuses 32-token x 64-dimension V tiles across GQA
heads. A stress audit found and fixed a missing maximum-to-denominator phase barrier shared by the
old parallel and new score kernels; 100 independent processes now reproduce both paths exactly.
Against the corrected baseline, split GQA improves the 32K full operator 2.67x and 1K SWA 1.74x.
Wrapped-ring differentials cover both 12B/31B shapes; authentic 2K logits retain argmax/top-10 set
with correlation 0.999793 and NMSE 4.11e-4. Three-repetition 2K 31B pp improves 33.29%, so this is
the default F32 attention path from eight live tokens onward.
The score pass now also reuses K across GQA heads. One workgroup stages K for an eight-head block
and 32-token chunk, a separate kernel performs exact global normalization, and the promoted shared-V
pass completes attention. This adds 27.41% at 1K full attention and 49.76% at 32K; the stable 4K
31B model gate improves pp 9.00% with tg unchanged outside noise. It is the current default F32
attention path; the prior split-score kernel remains selectable as a differential oracle.
Initial prompt attention now has a causal 32-row specialization. It stores the non-wrapping initial
K/V tile first, launches one grid over prompt rows and query heads, and keeps each row's complete
score vector in threadgroup memory for one exact global max/sum softmax. A 32-token prompt falls
from 8,012 to 3,152 dispatches. Interleaved same-process A/B improves 31B pp 6.85% (92.30 to 98.63
t/s), above the 4.63% noise floor; tg changes -0.23% inside variance.
Later non-wrapping tiles through 1,024 live tokens now batch the already-promoted exact split-GQA
algorithm across up to 32 query rows: one cooperative-K score grid, one global-softmax grid, and
one shared-V grid. The batched output is bit-identical to the rowwise operator regression. At
64/128/512 prompt tokens, throughput improves 7.51%/6.80%/4.26% in the native profiler while the
512-token dispatch count falls 61.57%. The formal interleaved 512-token 31B gate improves pp from
71.01 to 86.80 t/s (+22.24%, 6.91% noise floor); tg changes -1.11% inside variance.
The same route now extends from 1,024 to 4,096 live tokens for hax-sized project prompts. The
ceiling is explicit: 31B uses at most 16 MiB of batched score scratch, and longer/full-ring work
falls back instead of allocating for the 262K model limit. A >1K operator differential is
bit-identical to rowwise attention. The 2K/4K profiles improve 9.30%/9.17% and remove
29,760/89,280 dispatches. A five-pair interleaved 2K gate improves pp 6.50% (50.88 to 54.18 t/s),
above the 2.93% noise floor; tg changes +0.28% inside variance. Authentic 4K logits retain argmax
and top-10 set with correlation 0.999379 and NMSE 5.34e-4.
Full-tile stage profiling at hax position 2,496 then exposed a separate scheduling cliff: KV
attention consumed 62.20% of sampled GPU time because all 50 sliding-window layers reverted to
32 rowwise attention evaluations once their 1,024-token rings wrapped. The promoted path now runs
the batched split-GQA passes before committing a tile, reading older positions from the ring and
causal current-tile positions from pending F32 K/V activations, then commits the tile to the ring.
The wrapped CPU differential passes both model geometries. Fresh 31B 1,024/1,056/2,048-token
rowwise-oracle comparisons retain argmax and ordered top-10; at 2,048 tokens correlation is
0.999999793 and NMSE is 9.08e-7. On the authentic 2,557-token hax workload, a five-pair interleaved
same-process A/B improves pp from 48.42 to 59.53 t/s (+22.93%), above the 9.16% noise floor. A
64-token secondary interval changes tg +5.90% inside variance, with no regression. Three exact HTTP
replays return byte-identical `OK` from 2,557 prompt tokens. Evidence is in
`artifacts/2026-08-24-real-hax-prefill-31b`.
Prompt batch 16 was separately rejected on the same workload: it reduced pp from 54.52 to 31.05
t/s (-43.04%). Batch 32 remains the production tile.
The optional F16 KV route now uses the same batched split-GQA structure during prefill while
retaining its prior one-pass decode kernel. At 1,024 prompt tokens this single scheduling-axis
change improves pp from 49.14 to 82.22 t/s (+67.31%); tg changes -1.62%, inside baseline variance
and below the secondary regression limit. A direct F16 split-GQA decode variant was separately
rejected because its +3.2% 32K operator result became an -8.71% authentic generation regression.
The promoted hybrid is exposed as `--kv-type f16`, halves KV capacity cost, and remains opt-in;
F32 is the accuracy default because 2K F16 logits retain argmax but drift by 1.91e-3 NMSE. Published
int4 Gemma 4 31B results report compounding quality loss beyond roughly 950 tokens, so int4 KV is
closed at the research gate pending stronger long-context evidence. Evidence is in
`artifacts/2026-08-24-f16-split-gqa`.
## K4: scheduling and memory
- Precompute an immutable 48/60-layer dispatch schedule at load time.
- Allocate tensor and scratch arenas once. Generation performs no tensor-sized heap allocation;
the current Metal batch handle has one small host `calloc/free` per command buffer. Retained
production-path attribution bounds all non-GPU work at 0.102% of wall time and batch cleanup at
roughly 0.005 ms, so eliminating that control allocation is explicitly below the 3% speed floor.
- Use one command buffer per useful pipeline span, not per primitive operation.
- Keep Q4_0 weights mapped/read-only and KV/scratch in Metal shared buffers; test private-buffer
staging only if Instruments proves a win.
- Add prefix-cache support only after deterministic uncached continuations pass.
Status: token-span scheduling and the four-row Q4_0 decode kernel are promoted. The default top-k
40 sampler now uses bounded heap selection rather than a 262144-entry full sort, reducing measured
CPU selection latency from 26.52 ms to 0.96 ms (27.66x) with seeded full-sort parity. `top_k=0`
retains the exhaustive sort semantics.
Exact append-only prefix reuse is promoted for the serial hax server. The cache retains the prior
canonical request boundary rather than raw generated control tokens, snapshots every K/V slot that
a response can overwrite, and accepts only a strict token-for-token extension. At a 3.1K
conversation, second-turn latency falls from 74.14 to 1.057 seconds on average (70.14x); a wrapped
1,025-token rollback preserves argmax and has 3.29e-4 NMSE versus 3.56e-4 repeat-control NMSE.
The follow-up lazy token-major rollback removes the 128-token ceiling: checkpoint creation copies
nothing, evaluation protects slots on first write, and the backup grows only with the evaluated
tail. A wrapped 160-token ASan/UBSan gate preserves argmax with 7.47e-5 NMSE versus 1.85e-4
repeat-control NMSE. Three byte-identical HTTP second turns improve from 4.431 s full-prefill mean
to 2.391 s exact-hit mean (1.85x, -46.04%). Memory pressure safely drops reuse rather than generation.
Real hax completed a shell tool round with 2,533 reused and 34 appended tokens. Evidence is in
`artifacts/2026-08-24-exact-prefix-cache` and
`artifacts/2026-08-24-lazy-prefix-rollback`.
A paired half-block Q4_0 lane-layout candidate passed operator correctness but was rejected at the
31B FFN decode shape: 311.43 GB/s versus 335.93/344.42 GB/s surrounding four-row controls. The
experiment was removed. Decode's observed model-weight traffic and the row-four microbenchmark are
both approximately 338 GB/s, so the next required evidence is per-stage GPU timing for an authentic
one-token 31B layer. Only that trace may select the next axis: lossless fixed-shape Q4_0 repacking
or a projection/epilogue fusion. Requantization into a different native format is not permitted.
The same dataflow was screened once more against the 12B-only 3840-to-15360 FFN gate/up geometry
because its isolated operator result was positive. Full-model evidence overruled the microbenchmark:
five-run A/B/A decode changed 41.364 to 41.144 t/s (-0.53%), below the 3% floor and within elevated
candidate variance. Logits retained identical argmax and ordered top-10 with 1.69e-12 NMSE. The
12B selector and diagnostic kernel were removed; do not reintroduce this lane layout without a
materially different dataflow.
The batch-one Q4 benchmark fidelity gap is closed. Its persistent mode now rotates explicit matrix
slots, reports the total working set, and warms using the selected multi-dispatch command span.
Sustained 12B bandwidth loses 6.52% when moving from one 31.64 MiB matrix to a 253.13 MiB stream;
31B is unchanged because one 62.02 MiB projection already streams. A graph-like span raises 12B
streaming throughput 34.14%, and production already uses a token-wide command buffer. Use at least
250 MiB and graph-like span for future batch-one screens.
A four-row Q8 tied-logit candidate is complete and rejected. It passed CPU correctness and the
immediately preceding 1,154/1,154 reference MUL_MAT oracle, but regressed the clean 12B exact-shape
operator 10.23%. Any optimistic 31B stage gain was too small to clear 3% end to end because logits
are only 7.34% of sampled decode time. The candidate was removed before model integration.
The resulting one-token counter profile attributes 81.29% of sampled 31B GPU time to
projection-dominated groups, 6.04% to KV/attention, 7.34% to logits, and only 5.30% to explicit
norm/transform/activation groups. A lossless split-scale Q4_0 sidecar subsequently regressed to
324.21 GB/s versus 371.49/382.49 GB/s direct-layout controls and was removed. Single-token fusion
and repacking are therefore paused.
## K5: advanced candidates
Only after trace evidence:
- MTP/draft-model decoding and acceptance-rate accounting;
- packed Q4_0 layouts or function-constant kernels for fixed dimensions;
- fused Q4_0 projection + activation epilogues;
- concurrent CPU tokenization/server work with GPU execution;
- multimodal projector support (text-only hax serving remains the first completion target).
Status: the fixed-eight Q4_0 verifier, batched Q8 logits, wrapped-KV rollback, exact rejection
sampling, and official four-layer Gemma assistant are promoted. The full 12B drafter was rejected
by its measured acceptance/cost gate. A new MTP stage trace attributes only about 6.9 ms of a
33 ms seven-draft wall interval to sampled GPU arithmetic. The next bounded change therefore keeps
all seven recurrent drafts in one command buffer and selects top-k40/top-p tokens on device. Its
descending selector is independently adapted from the pinned MIT llama.cpp Metal TOP_K path after
445/445 backend-op cases passed on M4 Max. Full-vocabulary local differentials pass, all seven
greedy proposals match, and the interleaved three-repetition gate improves 34.262 to 31.764 ms
(-7.29%). The hax-compatible direct/device-MTP gate remains positive at -5.13% latency.
The fixed-eight verifier plus acceptance profile is complete. Exact CPU acceptance consumes only
about 1.87 ms per round, while target verification consumes about 159 ms and is 99.5% GPU time.
A 542-stage trace attributes 94.40% of sampled verifier time to Q4 projections. Tied logits are only
4.90 ms and attention is 1.34 ms, so device-side acceptance fusion is deferred. Doubling the fixed
Q4 kernel from four to eight SIMD groups regresses row throughput 8.74% and has been removed.
The two immediate four-group scheduling candidates are complete and rejected. A 128-output tile
regresses stable operator throughput 9.84%. A single-load dequantizer appears 11.3% faster when the
same 65 MiB matrix is repeated, but improves the authentic verifier only 0.30% inside variance;
both candidates were removed. The latter result exposes a benchmark-fidelity gap: cached repetition
is not representative of the full graph's unique weight stream.
The Q4 operator harness now rotates an explicit number of independently backed matrix slots and
reports total working-set size. At the 31B FFN shape, a four-slot 260 MiB stream initially matches
the stable single-matrix rate but later shows substantial resident-working-set variance. Cache reuse
alone therefore does not explain the single-load candidate's failure, and the streaming loop remains
a diagnostic rather than a promotion oracle.
A K=128 candidate subsequently reduced Q4 block rereads and barriers by four, but quadrupled
threadgroup scratch to 18 KiB and regressed 24.69%; it was removed. This agrees with the platform
boundary in Apple's current guide: K=128 is an M5 starting point, not an M4 recommendation. RIGEL's
M4 Max measurements also rule out Metal 4.1 FP8/tensor primitives as a shortcut because they execute
on the existing shader/simdgroup path and FP8 is emulated.
The scratch-free register-resident screen also failed: despite reading each Q4 block once, it reaches
only 10,886.60 rows/s versus 18,828.53 rows/s surrounding MMA controls (-42.18%). Repeated scalar
activation loads and FMAs are substantially worse than the staged simdgroup-matrix path at eight
tokens, so the production MMA structure remains the boundary.
The aligned-output specialization is complete and rejected. Removing the per-K-tile tail branch and
tail-store path reaches 19,068.84 rows/s versus 18,805.56 rows/s surrounding controls (+1.40%),
below the 3% gate and inside its 364.43 rows/s run standard deviation. The generic tail-safe kernel
and API were restored.
Packed-word Q4 loading is promoted. Replacing sixteen byte payload loads with eight aligned
`ushort` loads while preserving the exact centered-integer half multiply improves the rotating
operator 8.11% and authentic 31B verifier latency 4.52%; all 2,097,152 logits are bit-identical.
The faster algebraic half-scale rewrite was rejected for model-level numerical drift, demonstrating
that operation order is part of the correctness contract.
The batch-one transfer is complete and rejected. Exact packed-word unpack reaches 369.32 GB/s
against 367.92 GB/s surrounding controls (+0.38%) and is 1.73% below the stable post-control, so the
original byte loader remains production.
The 32-row prefill transfer is complete and promoted. Exact packed-word staging reaches 40,261.50
rows/s versus 38,165.60 rows/s surrounding controls (+5.49%) and improves authentic 31B prompt
throughput from 110.89 to 114.40 tokens/s (+3.16%). Generation is unchanged within noise (+0.29%),
and all 262,144 logits are bit-identical. Production therefore uses the packed loader only in the
64x32 prefill MMA and fixed-eight verifier; batch-one decode retains byte loads.
The refreshed authentic decode trace still assigns 36.06% of stable sampled time to the two FFN
gate/up projections, 19.67% to FFN down/output, and 16.24% to QKV projections. Exact gate/up launch
fusion was then screened and removed. It improves the isolated two-matrix operator from 384.83 to
408.96 GB/s (+6.27%) but regresses authentic 31B generation from 19.61 to 18.56 tokens/s (-5.33%).
Bit-identical logits confirm this is a scheduling/occupancy failure, not numerical drift. Separate
production dispatches remain faster in the full graph.
The next weight-pass amortization axis is fixed-eight tree verification. Linear MTP accepts only
1.92 proposals per C-code round and 2.08 per prose round, yet the target correction appears in the
assistant's top four on 9/11 and 7/12 mismatch rounds respectively. Agent-like repetitive output
already fully accepts 9/12 rounds, so the linear path remains the fast case. First build a diagnostic
tree verifier with explicit ancestry masks and positions, prove each node against an independent
sequential oracle, and only then benchmark adaptive tree scheduling in the server.
The ancestry-masked attention primitive, explicit depth-based RoPE, and complete diagnostic layer route
are implemented. A chain tree is bit-identical to the existing verifier for all 31B logits. All eight
branched nodes preserve the sequential oracle's top-1 token with cosine similarity >= 0.999999882;
normal and ASan suites pass 26/26. This remains diagnostic-only. The next single axis is an adaptive
linear-versus-tree server gate, followed by the required three-repetition 31B throughput comparison.
Do not gate MTP against direct decode using accepted-prefix averages alone. Four- and eight-round
window policies regressed representative code by 14.29% and 15.22%, respectively, and were removed.
The policy changes the rate at which the request reaches later easy regions, making its own signal
nonstationary. The next scheduler comparison must keep a multi-token main path in both arms and report
tokens advanced per target pass as well as wall-clock tokens/s.
The first main5-plus-three-branch scheduler is also rejected. Confidence gating rescued low-acceptance
prose branches, but regressed code and agent throughput by 6.41% and 7.54%, and hybrid greedy output
diverged at token 21 across repetitions. Do not switch verifier topology within a deterministic request
until numerical equivalence is strong enough to preserve near-tied argmax decisions. Keep the tree
primitives diagnostic and return to fixed-graph target/MTP cost reduction.
The fixed-eight direct-half scale load is rejected. It preserved every verifier logit bit-for-bit but
improved the rotating operator only 2.85%, then regressed authentic 31B verifier latency 4.10% in a
seven-repetition, same-process A/B/A comparison against surrounding controls. The explicit byte-to-
`ushort` reconstruction remains production. For small kernel deltas, use this alternating harness
pattern whenever independent processes show thermal or residency phase changes.
The promoted device MTP chain now has a 70-stage diagnostic trace. Its original coarse split is
52.34% across the four assistant layers, 29.01% in tied output/post-projection, 15.19% in top-k40,
and 3.47% in embedding/input projection. A 512-thread top-k block regresses that stage 14.43% and
adds seven dispatches, so production remains at 1024 threads. The finer split identifies the tied
1024-by-262144 vocabulary projection as more than 91% of its output group. Doubling the Q4 matvec
fanout to eight rows regresses the isolated stable operator 2.00%, that authentic stage 5.97%, and
sampled chain work 3.69%, so the experiment is removed. Keep the four-row schedule. The next single
axis removed generic block-loop and output-tail control at fixed K=1024 without increasing live row
accumulators. It regresses the stable operator 0.47% and improves only the two stable authentic stage
runs by 1.41%, below the promotion floor, so it too is removed. The next tied-output axis must change
memory handling. A diagnostic block-major SoA copy is now rejected: after its first cold control, a
seven-repetition same-process A/B/A gate reaches 432.63 GB/s versus 456.77 GB/s for surrounding
row-major controls (-5.28%), while adding 27.36 ms conversion and 144 MiB duplicate storage. The next
single axis preserved row-major GGUF bytes and compared the shared no-copy Metal resource with a
GPU-private copy. Seven paired GPU-timestamp repetitions improve only 0.36% on-device and 0.69% by
wall clock, for an 8.93 ms copy and 144 MiB allocation. Private storage is removed. The tied-output
projection is already close to the sustainable memory-streaming ceiling; further isolated layout,
constant-folding, or storage-mode changes need a stronger mechanism than generic bandwidth tuning.
The fused exact-shape tied projection plus initial top-k block sort is also rejected. Its 256
threadgroups of 1,024 threads retain the production arithmetic and match all 20,480 finite and
NaN/tie oracle pairs exactly, but stable same-process GPU timestamps regress the combined stage from
0.406014 to 0.431587 ms (+6.30%). The larger threadgroup reduces projection occupancy more than the
eliminated full-logits traffic and dispatch can repay, so the diagnostic kernel and API are removed.
Evidence is in `artifacts/2026-08-24-mtp-fused-projection-topk`.
The pair-valued frontier is promoted without changing projection topology. The initial sort writes
8-byte `(value, token)` entries and the eight merge levels plus sampler no longer gather full logits.
It matches 40,960 finite/NaN/tie ranked checks and all 162 sampled-output checks. Seven-repetition
same-process A/B/A GPU latency improves 18.54%; the three-repetition authentic 31B chain improves
top-k 15.41%, preserves all proposals, and changes wall latency by only +0.52%. Two scratch buffers
grow by 80 KiB total. Evidence is in `artifacts/2026-08-24-mtp-topk-pairs`.
Exact 40-pair truncation at every merge node is promoted. The eight levels now write 10,200 pairs
instead of 81,920, and scratch B shrinks from 80 to 40 KiB. Finite and NaN/tie outputs match at all
2..256-run tree depths. Isolated GPU latency improves 14.21%; an authoritative same-model seven-run
A/B/A gate improves authentic top-k 13.79%, total GPU time 2.34%, and wall time 2.32%, with identical
outputs. Evidence is in `artifacts/2026-08-24-mtp-topk-compact`.
The exact four-way merge is promoted. Twenty-one complete and missing-child shapes match the compact
binary tree for finite/NaN/tie frontiers and sampled outputs. Four levels replace eight, global merge
writes fall from 10,200 to 3,400 pairs, and scratch B shrinks from 40 to 20 KiB. Isolated GPU latency
improves 3.50%; the same-model seven-run 31B gate improves authentic top-k 3.27%, total GPU time 0.54%,
and wall time 0.48%, with identical outputs. Evidence is in
`artifacts/2026-08-24-mtp-topk-fourway`.
Terminal four-way merge/sampling fusion is rejected. It matches all 3,402 finite/NaN/tie sampled
outputs and removes one dispatch plus the final global frontier write, but stable isolated A/B/A
latency improves only 1.72% on-device and 1.82% by wall clock. The fused kernel and diagnostic API
were removed. Evidence is in `artifacts/2026-08-24-mtp-topk-terminal-fusion`.
Top-k is now bounded near 0.12 ms per draft. Use a refreshed authentic MTP trace to choose a larger
fixed-graph fusion or reuse axis rather than continuing dispatch-only sampler changes.
The refreshed stable trace places about 4.04 ms in assistant layers, 3.25 ms in output projection,
and 0.83 ms in top-k per seven-draft chain. Fusing post-attention RMSNorm, residual add, and FFN
RMSNorm was bit-exact over 131,072 floats and improved its isolated three-to-one dispatch benchmark
39.1%, but improved authentic layer time only 1.30%, total GPU 0.69%, and wall 0.69%. It is rejected
and removed. Evidence is in `artifacts/2026-08-24-mtp-post-attention-fusion`.
The assistant gate/up-plus-GELU fusion is also rejected and removed. It was bit-exact, eliminated
56 dispatches per seven-draft chain, and avoided roughly 96 KiB of F32 intermediate traffic per
layer, but the stable same-model 31B gate improved assistant layers only 1.25% and total sampled GPU
and wall latency 0.95%. The unchanged 9 MiB canonical Q4 weight stream dominates this saving, and
the combined accumulators add pressure. Evidence is in
`artifacts/2026-08-24-mtp-gate-up-gelu-fusion`.
Further assistant-layer work must remove arithmetic or weight traffic. First add a fine-grained
authentic layer trace; do not select another fusion from dispatch counts alone.
The original five-token fine trace is superseded for production attribution: the harness had
hardcoded F16 KV. With the harness corrected to the production F32 default, two stable 580-token
31B traces place 54.79% of assistant-layer time in attention; a third broad interference run is
retained as an outlier, and all proposals match. F16 attention is slower at the same context. A
four-head/two-token full-score geometry then passed local correctness but improved median dispatch
throughput only 0.33%, below the 3% screen, and was removed. Evidence is in
`artifacts/2026-08-24-mtp-layer-fine-profile-f32-580` and
`artifacts/2026-08-24-mtp-score-head-block-screen`.
The refined F32 trace splits attention into score, normalization, and value passes. Across three
stable 580-token runs they average 1.557/0.136/2.730 ms, making shared-V accumulation 34.41% of all
assistant-layer work. A four-KV-head GQA2 packing preserved bit-exact wrapped-ring outputs and
halved value workgroups, but regressed stable operator throughput about 26%; synchronization and
indexing outweighed SIMD utilization. It was removed. Evidence is in
`artifacts/2026-08-24-mtp-attention-substage-profile` and
`artifacts/2026-08-24-mtp-value-gqa2x4-screen`.
A one-KV-head GQA2 control then doubled only the value dimension tile from 64 to 128. It preserved
bit-exact outputs and halved workgroups without cross-KV indexing, yet regressed stable throughput
21.33%; the 17 KiB threadgroup footprint and four per-lane accumulators are too costly on M4 Max.
It was removed. Do not pursue workgroup-count reduction alone; the next value design must reduce
V traffic or synchronization. Evidence is in `artifacts/2026-08-24-mtp-value-gqa2-d128-screen`.
A token-tile 64 control held V traffic and output arithmetic fixed while halving value-kernel
barrier rounds. It remained bit-identical and improved stable wall/GPU operator throughput only
0.97%/1.27%; the 18 KiB threadgroup footprint absorbs most of the synchronization saving. It was
removed below the 3% floor. Evidence is in `artifacts/2026-08-24-mtp-value-token64-screen`.
Token tile 48 used 13.5 KiB and reduced barrier rounds by roughly one-third. It improved stable
wall/GPU throughput 1.99%/1.78%, better than tile 64 but still below promotion, and was removed.
The 32/48/64 curve closes token-tile tuning for this value dataflow. Evidence is in
`artifacts/2026-08-24-mtp-value-token48-screen`.
The different-dataflow simdgroup-matrix value screen is also rejected. It padded GQA2's two useful
rows to an 8-by-8 primitive and converted F32 weights/cache values to half before float accumulation.
Although CPU-relative wrapped-ring 12B/31B checks passed, stable wall/GPU throughput regressed
40.92%/45.57%. The diagnostic was removed before an oracle/model gate. Matrix attention should be
reconsidered only when the storage format or batch shape supplies enough useful rows to amortize
conversion and padding. Evidence is in `artifacts/2026-08-24-mtp-value-mma-screen`.
The first reduced-weight screen is rejected before model integration. An independently adapted
14-byte Q3_0 split-bitplane kernel passed its CPU/Metal differential, but regressed exact assistant
projection latency 71.76% at 1024-by-8192 and 56.61% at 8192-by-1024. Its 22.22% byte saving cannot
repay low/high bit extraction on M4 Max. The diagnostic kernel, API, tests, and benchmark switch were
removed. Evidence is in `artifacts/2026-08-24-mtp-q3-weight-screen`.
Do not pursue smaller bits as an end in themselves. The next weight-traffic candidate needs a layout
whose decode maps to native packed operations, or a mechanism that reuses one canonical Q4 weight
pass across multiple useful assistant states without changing request determinism.
Linear MTP K=16 is closed for the current assistant and verifier graph. A dedicated fixed-16 Q4
kernel improved equivalent projection throughput 35--38% and passed both local and adjacent Metal
correctness gates, but three authentic 31B repetitions reduced generated throughput from 28.44 to
22.44 t/s (-21.11%). Mean accepted prefix rose only from 4.5 to 5.67 while verifier latency rose
from 1.92 to 3.15 seconds per 12 rounds. The diagnostic kernel and K=16 APIs were removed. Revisit
longer lookahead only if the verifier avoids proportionally evaluating all target layers or measured
acceptance changes materially. Evidence is in `artifacts/2026-08-24-q4-verify16`.
The refreshed canonical 31B A/llama/A scorecard records native F32 at 53.66 pp and 20.27 direct tg
t/s versus pinned llama.cpp F16 at 31.16 pp and 23.52 tg t/s. Native wins short prefill by 72.21%,
llama.cpp wins direct decode by 16.05%, and native K=8 MTP at 28.44 t/s remains 20.92% faster than
the llama.cpp direct lane. A Q4 resource-specific encoder barrier screen was removed as inconclusive:
candidate and post-removal control were within 0.24%, while the system drifted 8.86%. Direct-decode
work now requires stable interleaved evidence; the fastest production lever remains K=8 MTP.
Evidence is in `artifacts/2026-08-24-current-31b-scorecard` and
`artifacts/2026-08-24-q4-resource-barrier`.
The one-command hax launcher now automatically selects a regular colocated
`mtp-<model filename>` assistant for F32 `gem agent` runs. Explicit `--mtp` and `--no-mtp` retain
full control, while `gem serve` remains explicit. A three-pair real-hax gate reduced complete process
latency from 4.097 to 3.317 seconds (-19.04%, 1.235x) with six byte-identical greedy outputs and
verified server cleanup. This closes the integration gap between the promoted K=8 engine path and
the preferred offline command. Evidence is in `artifacts/2026-08-24-agent-mtp-autodiscovery`.
Real-model cancellation is now a cooperative server invariant. SSE keepalive comments before/after
prompt evaluation and per generation round detect an abandoned hax stream even while Gemma emits no
visible text; `MSG_NOSIGNAL` prevents disconnects from killing the process. The authentic 31B MTP
gate recovered health in 1.003 seconds with the server alive. A same-binary three-pair hax gate found
no latency regression and preserved byte-identical output. A redundant poll-before-write candidate
was removed below the 3% threshold. Evidence is in `artifacts/2026-08-24-real-cancellation`.
The deployed 31B F32 context gate is complete at 32K, 64K, and 128K with K=8 MTP. All three
capacities reached readiness and returned three correct non-streaming OpenAI completions; the
steady exact-prefix request means were 1.346/1.350/1.351 seconds despite KV allocations increasing
from 7.05 to 23.15 GB. This gate also closed the server/probe contract gap by adding native
`stream=false` aggregation. A low-level ordered sweep showed a nominal 128K decode drop, but
bracketing 32K/128K reruns were bimodal independent of capacity, so no context-performance claim or
kernel change is promoted from those values. Evidence is in
`artifacts/2026-08-24-context-capacity-31b`.
True hax-visible TTFT is now a first-class gate through `scripts/server-stream-probe.sh`; transport
headers are reported separately. The promoted emit-before-followup schedule moves only already
sampled or target-verified tokens ahead of the next-logit evaluation. At 32K, same-process
interleaving improves 31B K=8 visible TTFT 11.42% and total latency 3.78%, while direct TTFT improves
8.18%. Output events and usage remain exact. A root-only early emission candidate was removed
because Gemma's initial channel-control pieces are intentionally invisible. Evidence is in
`artifacts/2026-08-24-streaming-ttft-31b`.
A 31B q8k8 diagnostic reuses K across eight query rows for both 512-dimension full attention and
wrapped 256-dimension SWA while retaining global normalization and shared-V passes. After a stress
gate exposed and fixed a missing SIMD-group visibility barrier, 200/200 independent processes pass
and exact operator score time improves 79.27%/82.62%. Post-fix 256/512 model means improve
4.44%/4.27%, but both remain inside large host/thermal noise floors. Do not make this path production
until a stable post-fix model gate passes; the scalar score route remains default. Evidence is in
`artifacts/2026-08-24-full-score-q8k8-31b`.
The 31B row4 long-attention diagnostic is now decisively rejected for production. Although its
isolated attention operator remains roughly 15--17% faster than row2, a corrected interleaved 32K
prefix-extension gate including warmup and final logits makes row4 5.54% slower, with all three
pairs favoring row2 and identical logits. Two-way sequence splitting, direct device K/V reads, and
a 512-thread geometry also regress their row4 operator baseline. The experimental variants were
removed; row2 remains production. Before selecting another long-prompt kernel axis, add authentic
long-prefill stage attribution that separates projections/FFN, RoPE/KV, SWA, global attention, and
final projection. Evidence is in `artifacts/2026-08-25-long-context-row4-followup`.
That attribution is now complete at a production 4K position. SWA attention consumes 53.35% of
sampled tile GPU time versus 6.55% for global attention; projections and FFN are secondary. Three
SWA reuse strategies were rejected despite 9--38% extension speedups: full 32-row batching, a new
online row-pair kernel, and partial one-/two-of-five batched routing all exceed the model-logit NMSE
limit after repeated 50-layer use. The online kernel passed its native operator and the 4,793-case
reference gate but was still numerically unsafe end to end, so it and every temporary selector were
removed. The next SWA design must preserve rowwise arithmetic order, not merely approximate its
operator output. Evidence is in `artifacts/2026-08-25-long-prefill-stage-attribution`.
A follow-up synchronization audit restored the missing shared-memory barrier between maximum
reduction and sum reduction in the batched normalizer. This is required by Metal and remains as a
correctness fix, but it did not change the full-model result: the batched SWA route was 15.50%
faster at the 4K boundary while reproducing the rejected correlation 0.998577 and NMSE 3.65e-3.
The diagnostic selector was removed. Before any further SWA kernel geometry work, build a
first-divergence harness that captures authentic per-layer attention intermediates and KV state.
Evidence is in `artifacts/2026-08-25-batched-swa-barrier`.
The first-divergence work has now converted that rejected path into a production win. The only
causal mismatch was the empty-cache tile: batching it introduces layer-zero NMSE 5.45e-13, which
changes layer-one K/V and then amplifies. Keeping tile zero rowwise makes all 128 tiles through 4K
commit bit-identical K/V. Subsequent SWA tiles use the three-dispatch batched scalar route without
encoder rotations. Shared-prefix A/B improves 4K by 24.29% and 32K by 14.57%, with repeated logits
bit-identical. The latest promoted stage samples reduce SWA attention by 60.56% versus the recorded
baseline, as attribution only. This route is now the default
31B/F32 configuration; next measure its cold 32K/64K/128K server TTFT and retain the separate
prefix-extension gates. Evidence is in `artifacts/2026-08-25-swa-first-divergence`.
The next SWA kernel axis is measurement-gated query-row reuse without dot-product reassociation.
First capture promoted scalar score/normalize/value timing at a wrapped 1,024-token window and an
Instruments Metal trace with occupancy, limiter, and bandwidth counters. If scalar Q*K still
dominates, test a score-only two-row tile: stage each K token once, assign SIMD groups across the
two causal rows and their two GQA query heads, and retain the current per-lane accumulation plus
`simd_sum` order for every dot product. Keep global normalization and P*V unchanged. Compare a
four-row tile only if counters show the two-row geometry is not resource- or occupancy-limited.
This is distinct from the rejected online row-pair kernel, which changed the complete softmax/value
arithmetic. TensorOps remains a separate M4 fallback experiment after this exact-order axis.
The two-row exact-order score axis is now promoted. It reduces isolated wrapped-SWA score time
38.91%, complete attention 24.30%, and authentic 32K extension time 7.14%, with bit-identical
logits. The production selector uses rowwise tile zero, scalar batched pre-wrap SWA, two-row score
reuse after wrap, and online row2 full attention. Reprofile this route before selecting another
kernel; the previous score/value balance has changed. Evidence is in
`artifacts/2026-08-25-swa-score-row2`.
The promoted operator profile now puts score at 1.45 ms and P*V at 1.38 ms, with normalization at
0.038 ms, so softmax-only fusion is closed and neither remaining half has enough isolated evidence
to choose by wall time alone. Replay the retained Metal trace in Instruments with occupancy,
limiter, utilization, and bandwidth counters after the 64K gate. Register/threadgroup pressure
closes larger row tiles; memory-read stalls with occupancy headroom admit an exact-order four-row K
reuse screen; dominant normalized-score/V traffic selects a value-side reuse screen. Change one of
those axes, not both.
For post-prefill generation, a Paperbridge audit of Open-TQ-Metal identifies a separate potential
128K axis: many-chunk sequence-parallel decode with a second-phase merge of partial online-softmax
state. Do not combine it with KV quantization; the same study reports Gemma 4 int4/int8 KV failure
near 1K tokens and lacks Gemma long-context quality evidence. The earlier rejected two-segment row4
candidate was a 32-row prefill geometry and does not settle single-row 128K occupancy. First add an
authentic populated-position decode-stage profile; proceed only if the ten full-attention layers
offer at least a 3% end-to-end ceiling.
The decode operator profile now selects that axis more narrowly. At exact 31B full-attention
geometry, P*V consumes 77.01% at 64K and 77.04% at 128K while sustaining only 49.0/48.0 GB/s;
scoring sustains about 170 GB/s and normalization is at most 1.23%. SWA remains bounded to 1,024
tokens. Implement no fused score rewrite yet. First run the checkpoint-amortized full-model 64K
decode profile. If ten full-layer KV-attention time clears the 3% whole-token ceiling, first screen
a full-layer-only 32-dimension value tile. This doubles the exact 31B dispatch from 32 to 64
threadgroups, halves threadgroup V storage, and preserves the existing per-dimension token
accumulation order with no partial scratch. If it fails to clear the gate, then screen
sequence-parallel value accumulation and an F32 partial-output reduction as a separate axis. Test
split counts independently; do not combine split count with a new dimension tile or KV
representation.
The first screen is retained through the operator gate. The populated 64K model baseline assigns
76.40% of sampled token time to KV attention, and the isolated global P*V component supplies an
estimated 45.5% whole-token ceiling. The 32-dimension value tile is bit-identical at an 8,193-token
tail case and improves the bracketed 64K value pass by about 29.7% and complete attention by about
23.2%. It passes 43/43 normal tests, 43/43 ASan/UBSan tests, and the adjacent 4,793/4,793 MTL0
FlashAttention oracle. Keep it behind selector 18 until a populated 64K model A/B clears the 3%
gate; do not start split-K while this exact-order candidate remains viable. That final gate now
passes: authentic 64K token GPU time improves 5.51% with 3.11% maximum CV, KV attention improves
13.62%, and 8,193-token model logits are byte-identical. Selector 18 is the default 31B/F32 route.
The stable 128K closing bracket also improves value accumulation 29.23% and complete attention
20.64%. A later populated 128K model profile established enough residual headroom to admit the
separate normalized P*V sequence-partition experiment below.
If that populated profile admits another value-side experiment, keep the promoted 32-dimension tile
and screen normalized P*V sequence partitions 2/4/8/16 as four independent candidates. Each writes
one F32 partial output per query-head dimension and uses a deterministic ascending-partition reduction;
reusable scratch is 128 KiB through 1 MiB. This changes accumulation association but not score
normalization or aggregate V bytes, so require the nonmultiple-tail differential and repeated authentic
model-logit gate before any full-model timing. Do not introduce KV quantization in the same comparison.
Current upstream evidence strengthens the narrow, fail-closed version of this experiment. MLX's
merged GQA-8 two-pass optimization gains locality by assigning contiguous token subchunks to SIMD
groups, but deliberately supports only head dimensions 64/128 because register pressure already
changes the winning heads-per-threadgroup there; Gemma 4's dimension 512 is not a transferable
shape. A separate MLX report localized deterministic bad logits to its two-pass routing boundary,
showing that operator throughput and ordinary short-context tests are insufficient for a split
route. Therefore Gem's split P*V screen must test positions immediately below/at/above its selector
threshold, a nonmultiple partition tail, every split count, and repeated authentic model logits.
Apple's guidance also says occupancy is diagnostic only when correlated with limiter and bandwidth
counters. Do not choose a split count merely because it launches more threadgroups; require a lower
P*V wall time and verify that the merge dispatch plus scratch traffic does not erase the gain.
The authentic populated 128K profile now opens this experiment. Production dim32 decode averages
416.378 ms/token with 99.43% stage coverage; KV attention is 333.533 ms (80.56% of sampled time).
The matching exact-shape receipt attributes 233.321 ms to the ten global-attention layers and
160.513 ms to their P*V passes, respectively 69.95% of authentic KV-attention time and a 38.55%
theoretical whole-token ceiling. A 7.78% P*V improvement is sufficient to clear the 3% model floor.
Screen splits 2/4/8/16 in isolation now; this is authorization to implement diagnostic candidates,
not promotion.
That screen is complete. Split16 launches 1,024 dim32 value workgroups and a deterministic
ascending-partition reduction using 1 MiB scratch. It reduces authentic populated token GPU time
from 229.678 to 177.717 ms at 64K (-22.62%) and from 416.378 to 304.832 ms at 128K (-26.79%);
128K KV attention falls from 333.533 to 226.145 ms (-32.20%). All candidate samples are faster
than all matching baseline samples. Below/at/above-threshold model gates, deterministic repeats,
the adjacent 4,793/4,793 oracle, and complete normal/sanitizer suites pass. Selector 19 is therefore
the promoted 31B/F32 default; selector 18 remains the exact dim32 rollback. The next decode-kernel
axis is the now-dominant score/K pass, not another P*V split-count sweep.
That next score/K axis already has a bounded selector-20 diagnostic: it changes only the staged K
tile from one token to eight, retains eight-head GQA reuse and split16 value accumulation, and is
not production. Its schema-2 runner now requires clean pre/post 64K and 128K host receipts and runs
native/adjacent correctness after the timing arms. Resolve it after the current BK128 prefill and
128K server gates; do not implement a fused split-KV alternative while this cheaper axis is open.
The subsequent global-prefill cycle closes three tempting axes. Direct Family-9 cache reads regress
31-35%, explicit staged q8k8 regresses about 68%, and row4's 14-15% synthetic operator gain becomes
a 0.693% regression in an authentic 64K shared-prefix model gate. Selector 21 is removed and
selector 19 remains production. The retained 128K Metal System Trace plus focused row2 `.gputrace`
must now name the limiter before another prefill kernel is implemented; do not infer promotion from
seeded KV or isolated operator measurements.
The trace-directed research cycle then replaced scalar row2 ownership with a BQ8/BK64
SIMD-matrix kernel. It improves isolated 4K/64K/128K attention by about 59%, the authentic 4K graph
by 7.12%, and an authentic evaluated 64K graph by 33.38% while clearing argmax, top-10,
correlation, NMSE, determinism, sanitizer, and backend-oracle gates. Selector 19 now uses this path
for eligible 31B/F32 global-prefill tiles; selector 22 retains the old row2 route for rollback.
Exact conditional rescaling was then rejected and removed after sub-3% seeded results. The current
one-axis diagnostic holds BQ8, F32 KV, exact online softmax, head ownership, and dispatch geometry
fixed while doubling only BK64 to BK128. It is correct and deterministic and is routed only through
diagnostic selector 23: its first screen was host-contaminated, and the adjacent DeepSeek llama.cpp oracle independently uses
Q8/BK64/eight SIMD groups for the same 512-wide Apple Metal shape. The frozen retry requires ten
warmups and ten balanced interleaved samples at 4K/64K/128K, complete pipeline/device resource
receipts, both long-context gains above 3% and noise, and stable pre/post host state. Do not open a
second geometry axis until BK128 is either advanced to an evaluated-graph gate or removed.
## Promotion gate
Use three repetitions and a single changed axis. Promote only changes that exceed observed variance
and improve the selected KPI by at least 3% without more than 2% regression in the other KPI,
deterministic correctness, memory ceiling, streaming TTFT, or cancellation behavior. Preserve
negative results in project memory.