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.