gem/docs/research/architecture.md

Gem native architecture

gem is a standalone, model-specific inference server. Production code owns the entire path from GGUF bytes to streamed tokens and does not link libllama, ggml, or launch llama-server. llama.cpp DSV4 is an MIT-licensed reference implementation and the initial correctness/performance oracle, not a runtime dependency.

Gemma 4 GGUF
    -> mmap + strict 12B/31B tensor manifest
    -> Gemma 4 tokenizer and canonical chat formatting
    -> fixed 48/60-layer execution schedule
    -> Metal Q4_0/Q8_0/F32 kernels + unified-memory KV cache
    -> sampler
    -> minimal OpenAI HTTP/SSE API
    -> hax

Ownership boundaries

• gguf.c: bounded little-endian parsing, memory mapping, tensor extents, and recognized storage types. The production model contract currently accepts only the published layout: Q4_0 matrices, a Q8_0 token table, and F32 normalization data.
• model.c: exact dense Gemma 4 profiles and tensor names/shapes. It rejects other architectures, sizes, attention patterns, and unexpected tensor counts.
• weights.c: zero-copy pointers from the validated mapping into the fixed layer contract. Full layers bind K-as-V and proportional-RoPE factors explicitly.
• mtp.c: strict matching 12B/31B gemma4-assistant manifests and zero-copy bindings for the four-layer, 1,024-wide recurrent drafter. Assistant layers contain Q/O/FFN weights only and read the target's final SWA/full KV cache views.
• tokenizer: Gemma 4 vocabulary metadata, byte fallback, special tokens, and detokenization. It must match the reference token IDs before graph work is accepted.
• runtime: an explicit layer loop, scratch allocation plan, SWA/global KV views, sampling state, and cancellation. There is no general graph IR in the first engine.
• Metal: custom MSL kernels plus a small Objective-C bridge for Metal objects. The public engine and server APIs remain C. The full synthetic layer—Q/K/V, normalized RoPE, ring/global GQA attention, output projection, GELU-gated FFN, and residual epilogues—passes the scalar oracle. The page-rounded GGUF mapping is exposed as one persistent no-copy Metal buffer, tensors use validated offsets, and activations use one reusable workspace. The default scheduler spans a complete token across all layers. Prompt evaluation tiles up to 32 tokens through each layer, batches Q4 projections on Apple SIMD-matrix units, and preserves causal semantics by ordering RoPE, KV insertion, and attention inside the command buffer. RoPE derives all Q/K prompt positions in one grid per tensor. For the initial non-wrapping tile, all K/V rows are inserted first and one causal grid computes every prompt row with complete per-row scores in threadgroup memory; later tiles through 4,096 live tokens batch the exact split-GQA score, global-normalization, and shared-value passes across all rows in three dispatches. Once a sliding-window ring wraps, those passes read prior positions from the ring and causal positions in the current tile directly from pending K/V activations; the scheduler commits the tile to the ring only after attention. The 31B/F32 production selector deliberately keeps the empty-cache tile rowwise and uses this batched route from the second tile onward. After the 1,024-token SWA ring wraps, its score pass shares staged K across two adjacent query rows without changing scalar dot-product order. A 4,096-token lockstep scan verifies bit-identical committed K/V state against the rowwise oracle. The rowwise route remains selectable as a correctness oracle. The 31B batched score arena is capped at 16 MiB; longer full-attention work falls back rather than scaling scratch with the 262K maximum context. A diagnostic selector can replace scalar score generation on aligned 32-row 31B F32 tiles with fixed 8x8 SIMD-group matrices that reuse K across eight query rows. It supports full attention and wrapped SWA while feeding the same global normalization and shared-value passes, but remains unpromoted pending a stable post-fix model gate. Decode retains the dedicated four-output-row matrix-vector path. At eight live tokens, F32 attention switches from the tiny-context SIMD kernel to an exact split-GQA path: staged K score generation, global normalization, and 32-token x 64-dimension value tiles shared across query heads belonging to the same KV head. The unfused path remains a correctness and benchmark oracle. The cache remains F32 by default. The explicit --kv-type f16 profile halves KV memory, uses batched split-GQA attention for prefill, and retains the one-pass F16 decode kernel. Its 1,024-token prompt gate passed the model-level speed threshold, but it remains opt-in because its authentic 2K logit distribution has measurable drift from F32. The separate opt-in --kv-type f16-full layout keeps all 1,024-token SWA layers in F32 while storing only the context-growing full-attention layers in F16. At 128K this pins KV allocation to 12,415,139,840 bytes for 31B and 2,818,572,288 bytes for 12B; runtime-plan tests enforce both envelopes. It is a promoted opt-in capacity mode, not the default on the 128 GiB target because F32 remains the fastest absolute path. Both sizes route post-4K full-attention tiles through their shape-specific online F16 kernels while retaining F32 SWA. Repeated 2K/4K/4,128-token logits and interleaved 8K model gates cover that routing; 64K/128K cold and prefix-extension server timing remains an endurance gate rather than a prerequisite for using the explicit option.
• profiling: gem-decode-profile uses the M4's stage-boundary timestamp counter set. With one row it profiles production decode; --rows 2..32 profiles the corresponding production prefill tile at a selected position. It builds that prefix once and restores an exact KV checkpoint after the warmup and every sample, so 64K/128K diagnostics do not repeat prompt ingestion. It retains one command buffer but creates sampled compute passes for embedding, nine explicit stage groups per layer, and logits. This is a diagnostic path only. The ordinary inference API continues to encode the complete token without counter buffers or extra encoder boundaries. A real 31B parity gate confirms that sampled and ordinary one-token routes produce bit-identical logits. gem-decode-attention-profile separately exercises exact 12B/31B full/SWA geometry and reports the existing score, normalization, and value counter stages. It uses synthetic finite buffers for phase attribution and is never model-level promotion evidence.
• 31B/F32 global prefill after 4,096 live tokens uses a one-head BQ8/BK64 online-attention workgroup: eight query rows share each 64-token K/V tile, SIMD-group matrix operations produce the score blocks, and exact online softmax bounds scratch independently of context length. The pipeline uses 18,496 bytes of threadgroup memory and is production selector 19; selector 22 keeps the older row2 kernel as a rollback oracle. A BQ8/BK128 pipeline is compiled for controlled experiments only. It uses 20,544 bytes of threadgroup memory. Diagnostic selector 23 routes the otherwise-identical selector-19 graph to BK128 for evaluated-prefix/logit gates; no production selector uses it, and it cannot become production without the benchmark and correctness gates.
• 31B/F32 global decode at 8,192 live tokens or more uses a 32-dimension value tile and partitions its normalized P*V scan into 16 contiguous token ranges. The 1,024 first-pass workgroups write 1 MiB of reusable F32 partial scratch; a second dispatch reduces partitions in ascending order. Selector 18 remains the unsplit diagnostic rollback and selector 19 is the production default. Prefill, sliding attention, 12B, F16 modes, and short contexts retain their prior kernels.
• server: loopback HTTP with GET /health, GET /v1/models, and streamed POST /v1/chat/completions. SSE chunks use the OpenAI chat-completion delta shape and terminate with [DONE]. It also exposes llama.cpp-compatible /props metadata for hax probing, transforms OpenAI message/tool JSON into Gemma 4's native control-token grammar, and maps generated thought, text, and function-call blocks back into hax's SSE dialect. Hax's provider selector is still named llama.cpp, but gem agent redirects that adapter to this native endpoint; the label creates no llama.cpp runtime dependency. gem agent owns readiness and cleanup. With F32 KV, it also discovers a regular colocated mtp-<model filename> assistant unless --mtp overrides the path or --no-mtp selects the direct diagnostic route. gem serve keeps MTP selection explicit. Between serial requests, the server retains the exact tokenized request-prompt state. A later prompt reuses it only when it strictly extends every cached token; mismatches and identical replays reset normally. Response generation establishes a zero-copy checkpoint, then lazily snapshots each token-major set of K/V slots immediately before its first write. The backup grows geometrically through the remaining context, restores the canonical request boundary after streaming, and therefore has no fixed response-token ceiling. Allocation failure abandons only that cache opportunity and still serves the response. This avoids replaying multi-thousand-token hax conversations without approximate matching or unsafe ring rewinds. With --mtp ASSISTANT.gguf, the same native process runs lossless K=8 speculative rounds: one target-sampled root, seven recurrent assistant proposals, batched target verification, exact rejection/residual sampling, cache rewind, and bonus sampling. Before a verifier can wrap SWA, only the eight potentially overwritten slots are snapshotted per sliding layer; rejection restores the unaccepted tail. With top-k 40 and no repetition penalty, the seven recurrent graphs and nucleus draws stay in one Metal command buffer: an indirect Q8 embedding consumes each sampled device token and the drafter's post-projection writes directly into the next recurrent input. Sparse probability rows are copied once for exact target acceptance. Other sampler settings and short remaining outputs retain the CPU-draft/direct fallbacks. Protocol-valid SSE heartbeat comments are sent after every already-completed prompt batch and at each generation-round boundary. They are content-free cooperative cancellation points: a failed write aborts evaluation through the same reset/prefix-checkpoint cleanup paths, and MSG_NOSIGNAL prevents a disconnected client from terminating the process. One already-submitted prompt batch or speculative round completes before the next heartbeat. Prompts of at least 4,096 tokens additionally log batch-boundary progress, throughput, and ETA every 4,096 evaluated tokens. Requests with stream=false use the same native generator and aggregate text, reasoning, tool calls, finish reason, and usage into one OpenAI-compatible chat.completion JSON response. This compatibility path does not alter the hax streaming hot path. Once direct sampling or speculative verification makes a token immutable, streaming emits it before evaluating that token to prepare the next logits. This preserves exact generation while removing one target-evaluation interval from visible TTFT and avoids unused final-token work at EOS or the request limit.

The disabled src/gem.c subprocess launcher is migration scaffolding. New engine behavior lives in gem_engine; the production gem serve and gem agent paths never execute llama-server.

Fixed model profiles

Property	12B	31B
Layers	48	60
Embedding width	3840	5376
Query heads	16	32
SWA KV heads / head width	8 / 256	16 / 256
Full KV heads / head width	1 / 512	4 / 512
FFN width	15360	21504
Pattern	5 SWA + 1 full	5 SWA + 1 full
SWA / maximum context	1024 / 262144	1024 / 262144
Dense Q4_0 tensor count	667	833

Full-attention layers reuse the K projection as the V input and therefore omit attn_v.weight, but their cached vectors are not aliases: K receives the learned key norm while V receives unweighted RMS normalization. SWA layers have separate K and V projections. Dense 12B/31B set num_kv_shared_layers to zero internally, so trunk cross-layer shared-KV optimizations for Gemma 4 MoE variants do not apply. The optional assistant is a separate contract: its three SWA layers deliberately read target layer N-2 and its full layer reads target layer N-1.

Correctness ladder

Each layer must pass before the next becomes a performance target:

1. malformed/synthetic GGUF tests and authentic sparse-header validation;
2. tokenizer golden vectors against the reference converter/runtime;
3. scalar CPU operator tests, then a tiny-model/token-logit differential harness;
4. Metal operator comparisons against the CPU reference with absolute/relative tolerances;
5. deterministic first-token and short-generation comparisons on 12B and 31B;
6. hax discovery/request/SSE integration, real-model tool-call streaming, and cancellation tests;
7. three-run benchmark promotion gate on the physical M4 Max.

M4 Max operating assumptions

The primary artifact is Q4_0 because it has a simple 32-element/18-byte block, is published for both targets, and maps cleanly to a bandwidth-oriented Metal matrix-vector kernel. Q8_0 is retained for embeddings and as a quality/baseline option. Decode optimization comes first; prefill receives separate matrix-matrix and attention kernels. A packed sidecar is permitted only when a trace proves that direct GGUF layout materially limits throughput and it remains reproducibly derived from the canonical GGUF.

The default hax context should begin below the architectural maximum (32K is the initial capacity point) so the model, KV cache, scratch buffers, hax, and YAMS coexist with headroom. Context capacity and token latency are separate benchmark axes.