gem/docs/research/fused-split-kv-design.md

Parked D512 fused split-KV decode design

Status: research/design only. Do not add a selector or Metal pipeline until the BK128 prefill decision, current-production 128K server acceptance, and selector-20 score-tile8 decision are closed in that order.

The selector-independent CPU oracle prerequisite is now implemented in src/ops.c. It constructs online-softmax partials, merges empty and nonempty states with the exact rescaling identity, and finalizes the numerator. Unit coverage compares 1/2/3/4/7/16-way partitions over a 13-token, 17-dimension non-power-of-two fixture; the 16-way case necessarily includes empty partitions. Invalid finite-state combinations and scratch-size overflow fail closed. This code is not called by the production graph and does not add a Metal pipeline.

gem-online-merge-check wraps that oracle with an independent double-precision direct-softmax reference. Fixed P=32/DV=512 audits at 8,191, 8,192, 8,193, 65,536, and 131,072 tokens all produce finite, repeat-deterministic output. Correlation is at least 0.999999999999, NMSE is at most 1.42e-11, and maximum absolute error is below 1.64e-7. analyze-online-merge-boundary.sh binds those exact positions, geometry, thresholds, seeds, hashes, and raw-input digests. This discharges only the CPU algebra prerequisite; the future Metal candidate must still pass the same lengths plus greedy model continuations against forced selector 19.

Fixed production shape

This design is only for the ten full-attention layers of Gemma 4 31B at batch-one decode and an 8K-or-longer prefix:

• 32 query heads, four KV heads, eight query heads per KV head;
• symmetric DK=DV=512, F32 KV cache, no sliding window;
• selector 19's graph, Q4 weights, split16 value route, sampler, and server behavior remain the baseline;
• exact dense attention and deterministic execution remain mandatory.

The production "dim32" name is a value workgroup dimension tile. It does not mean the value head has width 32. The authentic GGUF contract requires both full-attention key and value lengths to be 512 (src/model.c).

Measured baseline

Selector 19 currently uses three attention stages plus the split16 value reduction:

1. The GQA score kernel stages each K row once for all eight query heads and writes 32*live F32 scores: 8 MiB at 64K and 16 MiB at 128K.
2. One normalization threadgroup per query head performs the full-row max, exponentiation, and sum.
3. The dim32/split16 value kernel divides each 512-wide output into sixteen dimension tiles and the sequence into sixteen partitions. It dispatches 1,024 workgroups and writes 1 MiB of partials.
4. A final elementwise kernel reduces the sixteen value partials.

Retained exact-shape profiles put the 64K score, normalize, and value stages at approximately 3.12 ms, 0.11 ms, and 2.65 ms. At 128K, stable samples are approximately 6.20 ms, 0.33 ms, and 5.38 ms. Score/K is therefore about half the remaining attention time, but a candidate still needs an authentic whole-token gain; operator time alone is not a promotion claim.

The reproducible fused-split-kv-ceiling.json attribution compares those exact-shape profiles with the retained selector-19 authentic stage profiles. The ten global-attention calls account for a theoretical maximum of 33.11% of whole-token GPU time at 64K and 43.73% at 128K; score alone accounts for 17.54% and 22.02%. Clearing the 3% whole-token admission floor therefore requires at least a 9.06%/6.86% reduction within global attention at 64K/128K. This admits one diagnostic but cannot predict its gain because compulsory K/V reads, dot products, exponentials, and reduction remain.

The minimum logical intermediate traffic explains why fusion is worth one bounded experiment. Let S=32*live*4 be score-scratch bytes (8 MiB at 64K, 16 MiB at 128K). The score stage writes S. Normalization makes one max read, one read/write exponentiation pass, and one read/write scaling pass (5S). Each of the sixteen dim32 value groups then reads the complete normalized score rows for its output-dimension tile (16S). The route therefore issues 22S, or 176/352 MiB at 64K/128K, of logical score traffic in addition to reading K and V once. GPU caches may satisfy part of it, so this is not a device-bandwidth claim; the candidate must confirm the avoided traffic with Metal counters. Its partition partials add only bounded P-dependent traffic.

Candidate ownership and storage

The first candidate is a two-pass exact online-softmax kernel with one first-pass threadgroup per (kv_head, sequence_partition) and eight SIMD groups per threadgroup. Each SIMD group owns one of the eight query heads sharing that KV head.

For each sequence tile:

1. The 256 threads cooperatively stage the tile's K and V rows once in threadgroup memory.
2. Each SIMD group computes its query head's score using simd_sum.
3. Each SIMD group updates its own exact online-softmax (M,S,O[512]) state in sequence order.
4. The threadgroup writes one (M,S,O) partial per query head and partition.

Each lane owns 16 output dimensions, so one query and one output accumulator require 32 F32 values per thread before scalar state and temporaries. Start with a four-token tile: K and V staging uses 4*512*2*4 = 16,384 bytes. An eight-token tile consumes the full public 32,768-byte limit before any other threadgroup state and is not the first candidate. Only an Xcode GPU trace may establish whether the D512 register and 16 KiB threadgroup footprint leave useful occupancy.

The independently released Open-TQ-Metal Gemma kernel provides compilation-level evidence for the same D512 per-lane register shape: each lane holds 16 query and 16 output floats. Its ownership is otherwise unsuitable here. It launches one threadgroup per query head, maps that head back to a KV head, rereads shared GQA K/V for every query head, and serializes a 16-SIMD-group output reduction through barriers. Gem's candidate instead uses eight SIMD groups for the eight query heads of one KV head and must prove that its additional 16 KiB tile staging does not destroy occupancy.

The pinned adjacent DSV4 vector kernel also proves why Gem cannot preserve GQA ownership by merely replicating its threadgroup layout. At symmetric D512, cache tile 32, and one SIMD group, its exact scratch formula is (pad(DK,128) + 4*C + 2*pad(DV,128))*2 = 3,328 bytes per query head: shared Q, score/mask state, and an F32 output accumulator. Eight heads would require 26,624 bytes before K/V staging. Adding Gem's tile4 shared K/V allocation would total 43,008 bytes, exceeding the M4 Max's 32,768-byte limit. The first candidate must therefore keep Q and O lane-local in registers and use threadgroup memory only for the shared K/V tile and small coordination state. If that register shape lowers the compiled pipeline below 256 launchable threads, the fixed candidate fails its resource gate; it must not silently fall back to a four-head layout that rereads K/V or changes the one-axis experiment.

The second pass merges partition states with the exact identity M=max(Ma,Mb), S=Sa*exp(Ma-M)+Sb*exp(Mb-M), and the corresponding rescaled output numerator. It then writes O/S. The partial layout is partition-major and contiguous in output dimension so each lane reads its 16 dimensions coalesced. At P partitions, output partials consume P*32*512*4 bytes: 2, 4, 8, and 16 MiB for P=32, 64, 128, and 256 respectively, plus small M/S arrays. This is bounded by context-independent P and replaces the linear score scratch on the candidate route.

gem_attention_online_scratch_bytes fixes the complete byte contract before Metal allocation. For P=32 it reports 4,096 bytes each for maxima and denominators, 2,097,152 numerator bytes, and 2,105,344 bytes total. Larger partition counts scale this entire layout linearly and all products are checked for size_t overflow.

The future compiled resource receipt is also fixed before implementation. The first pass must report SIMD width 32, at least 256 launchable threads, exactly 256 dispatched threads/eight SIMD groups, and exactly 16,384 dynamic threadgroup bytes for tile4 K+V staging. The reducer must report SIMD width 32, at least 256 launchable threads, exactly 256 dispatched threads, and no dynamic threadgroup allocation. Static plus dynamic memory for either pipeline must fit the M4 Max's 32 KiB limit. analyze-fused-split-kv-resources.sh additionally binds D512/P32/F32/no-window geometry, the 2,105,344-byte scratch contract, device family, and source/binary digests. Passing this receipt only proves the pipeline is compilable and launchable; it cannot substitute for an occupancy trace, correctness, or timing.

The subsequent correctness artifact is likewise defined before implementation. analyze-fused-split-kv-correctness.sh consumes the digested CPU gate and compiled-resource analysis, five raw operator differentials at 8,191/8,192/8,193/65,536/131,072 tokens, and three raw evaluated-prompt model differentials at the routing boundary. Operator records must preserve the fixed D512/P32/tile4/F32 geometry, use selector 19 as baseline, be finite and repeat-deterministic, and clear correlation 0.9999/NMSE 1e-4. Model records additionally require 60-layer 31B identity, unchanged production selection, matching argmax and ordered top-10, repeat-deterministic candidate logits, and at least eight identical greedy continuation tokens. Every raw file is bound by digest. This gate cannot report promotion eligibility; clean-host timing and the adjacent backend-op oracle still follow it.

One-axis experiments

Do not mix these experiments:

1. Correctness-only D512 fused split-KV implementation at fixed tile4 and P=32.
2. Sequence partitions P=32 versus 64, then the winner versus 128; admit 256 only if a trace shows insufficient first-pass parallelism. Each comparison is a separate one-axis ABBA run.
3. Only after selecting P, compare tile4 with another tile size under a separate gate.

Direct V reads, F16/quantized KV, paged KV, sparse attention, Metal 4 TensorOps, and multi-request scheduling are separate axes and excluded from this cycle.

Fail-closed evidence

Before authentic model timing, require:

• finite seeded CPU differential tests at exactly 8,191, 8,192, and 8,193 live tokens, plus 64K and 128K tails; the CPU-only algebra portion of all five lengths is complete, while Metal differentials remain;
• a forced-selector-19 comparison over the same boundary inputs so changing dispatch routes is numerically invisible, including matching greedy tokens across a multi-token continuation;
• repeat determinism, finite output, matching argmax/top-10, correlation at least 0.9999, and NMSE no greater than 1e-4;
• exact partial-merge tests covering empty/tail partitions and non-power-of-two live lengths;
• buffer overflow/range tests and release plus ASan/UBSan suites;
• a profile receipt recording pipeline thread width, maximum threads, static/dynamic threadgroup memory, raw GPU timestamps, and an Xcode occupancy/limiter trace; the fail-closed compiled-resource receipt analyzer is ready, while an authentic receipt requires the future pipeline;
• three interleaved repetitions at 64K and 128K with the frozen clean-host policy;
• authentic deterministic logits and whole-token timing before any selector change;
• the correctness analyzer contract for these operator/model records is ready, while authentic raw records require the future Metal pipeline;
• the adjacent backend-op correctness oracle after a kernel change.

The candidate must beat selector 19 by more than 3% and observed variation at both long contexts, without a short-context regression beyond the protocol. Passing an operator screen only admits the authentic model gate; it never changes production directly.

The ceiling screen is regenerated with scripts/analyze-fused-split-kv-ceiling.sh; it validates the selector, exact D512 geometry, context, stage totals, repetitions, sample coverage, and input digests before making the non-promotional attribution.

Primary design references

• MLX D64/D128 duplication-free GQA and two-pass reduction at d73eb752ef2e6288fd95b032c0bff0a15a4a9e93.
• MLX issue 4022, which localizes a deterministic GQA correctness failure to a two-pass routing boundary and motivates Gem's explicit boundary and continuation canaries.
• Open-TQ-Metal's D512 Gemma specialization at mutable-state-inc/gemma4metal@0f09466b7fde772a4876bf7bee3ccdeb34313304, used only as register-shape evidence. Its one-query-head ownership and short-context int4 results are not a performance or correctness oracle for this design.
• Adjacent DeepSeek-V4 llama.cpp vector online-attention/reduction oracle at d8df12ebc4592b55dcecf97a32953623d031efdc.
• Apple Xcode occupancy, memory-bandwidth, and limiter counters. Static resource limits alone do not establish resident workgroup count.