688 lines
51 KiB
Markdown
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.
|