188 lines
14 KiB
Markdown
188 lines
14 KiB
Markdown
# 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.
|
|
|
|
```text
|
|
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.
|