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

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)

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