gem/docs/research/kernel-roadmap.md

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 QK 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 PV 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 PV 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 PV 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 PV 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 PV 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 PV passes, respectively 69.95% of authentic KV-attention time and a 38.55% theoretical whole-token ceiling. A 7.78% PV 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.