Trede98/gemma-4-E2B-coreml

gemma4-e2b-ctx8192-int4-mf.mlpackage — integration guide

A self-contained Core ML port of google/gemma-4-E2B-it for on-device inference. Single .mlpackage, no external runtime.

What it is

Base	google/gemma-4-E2B-it (Gemma 4 / 3n, 35 layers, 1536 hidden)
Format	Stateful Core ML multifunction .mlpackage
Size	~2.3 GB (int4 palettized, per-grouped-channel g32)
Context	8192 tokens
Functions	prefill (T=128) and decode (T=1), sharing weights + KV state
Min OS	iOS 18 (MLState)
Verified	100% greedy parity (fp16) / coherent (int4) vs the HF reference

Functions

Both functions share one weight set (deduplicated at merge) and the same MLState KV cache. Load each with MLModel(function_name:); create the state once and pass it to both.

• prefill — input_ids [1, 128]. Ingest the prompt 128 tokens per call (chunk longer prompts; the cache carries context across chunks). Right-pad the final partial chunk. Returns all-T logits [1, 128, vocab] — read the row at the last real token for the first generated token. Measured ~190–260 tok/s.
• decode — input_ids [1, 1]. One token per step. Returns [1, 1, vocab].

Inputs (per call)

Name	Shape	Notes
input_ids	[1, T] int32	token ids (T = 128 prefill / 1 decode)
position_ids	[1, T] int32	absolute positions
causal_mask_full	[1, 1, T, 8192] fp16	0 allowed, −65500 masked; key ≤ query_pos
causal_mask_sliding	[1, 1, T, 512] fp16	windowed causal over the 512-slot ring
place_full	[T, 8192] fp16	one-hot: token t → full-cache column = its abs position
place_sliding	[T, 512] fp16	one-hot: token t → sliding column = (abs pos % 512)

place_* are the KV-cache write addresses — one-hot rows telling the graph which cache column each token writes to. They replace a dynamic slice index (which Core ML's tracer would freeze to a constant). The Python reference for building all four tensors is build_inputs() in stateful_wrapper.py — port that verbatim to Swift.

States (MLState, all fp16)

15 key/value pairs, key_cache_{i} / value_cache_{i} for i in 0..14 (the layers before the KV-sharing point; later layers reuse these in-graph):

• sliding layers: [1, 512, 256]
• full-attention layers (4, 9, 14): [1, 8192, 512]

Create with makeState(); reset (zero) between independent notes.

Driving one generation (pseudocode)

state = model.makeState()
// prefill in 128-token chunks
for chunk in prompt.chunks(128):
    ids, pos = chunk (right-pad last chunk to 128)
    logits = prefill.predict(ids, pos, masks, places, state)
next = argmax(logits[lastRealToken])
// decode
for step in 0..<maxNewTokens:
    if next in {1, 106} { break }        // <eos> / <end_of_turn>
    emit(next)
    logits = decode.predict([next], [pos], masks, places, state)
    next = argmax(logits[0])

Apply the Gemma chat template (from hf_model/tokenizer_config.json) when building the prompt; use swift-transformers' tokenizer for encode/decode.

Regenerate

hf download google/gemma-4-E2B-it --local-dir hf_model
.venv/bin/python build_multifunction.py --ctx 8192 --prefill-t 128 --quantize int4 \
    --out output/gemma4-e2b-ctx8192-int4-mf.mlpackage
CPUONLY=1 .venv/bin/python mf_parity.py \
    --pkg output/gemma4-e2b-ctx8192-int4-mf.mlpackage --threshold 0.6