Wiring CoreML to on-device LLMs: KV-cache, ANE, and the token budget for iPhone 16

Meta description: CoreML’s ML Program format for on-device transformer inference — ANE scheduling, 4-bit quantization tradeoffs, and the KV-cache formula that keeps you under iOS memory limits.

---

TL;DR

Running a quantized LLM on-device with CoreML requires understanding three interlocking constraints: how the Apple Neural Engine schedules transformer layers, which quantization format keeps weights under 2GB, and how KV-cache growth caps your usable context window before jetsam kills your process. Get any one wrong and you’re debugging OOM crashes in production.

---

Why on-device inference is a hardware scheduling problem

Apple’s unified memory architecture — where the CPU, GPU, and Neural Engine share the same physical memory pool — is both an advantage and a constraint. The advantage: no PCIe bandwidth bottleneck between host and device memory. The constraint: everything competes for the same pool, and iOS will reclaim it without mercy.

CoreML’s ML Program format (.mlpackage with a model.mlmodel specification) is the abstraction layer between your Swift code and that hardware. For transformer inference, the ML Program compiler partitions the compute graph across available hardware. The mistake most teams make here: assuming the ANE and GPU handle attention and feed-forward layers the same way. They don’t.

---

ANE vs. GPU: attention and feed-forward are not the same

The Apple Neural Engine is optimized for fixed-shape, low-precision matrix multiply. That maps well to feed-forward layers (FFN), where weight dimensions are static and the operations are embarrassingly parallel. Attention is a different story.

Layer Type	Preferred Backend	Reason
Feed-forward (FFN)	ANE	Fixed shapes, 4-bit weight load, high throughput
Multi-head attention (MHA)	GPU	Dynamic sequence length, flexible memory access
Embedding lookup	CPU	Small, irregular access pattern
Layer norm / softmax	ANE or GPU	Compiler decides based on input rank

When you export a model with coremltools, annotating compute units with ComputeUnit.ALL lets the compiler make this scheduling decision. Forcing ANE_ONLY on attention blocks is a common mistake — it triggers silent fallback to CPU and destroys your latency budget.

---

Quantization: 4-bit palettization vs. block-wise float16

To fit under the 2GB weight threshold that keeps iOS from aggressively reclaiming memory, you have two practical options.

4-bit palettization maps each weight to an index in a learned codebook. This is what coremltools.optimize.coreml.palettize_weights produces. Compression is roughly 4x over float16, and the ANE loads palettized weights natively on A17 Pro and M-series chips.

Block-wise float16 (per-channel or per-group quantization) retains float16 precision per block of N weights. Less aggressive compression, but better for layers sensitive to weight distribution — specifically the first and last projection layers.

// Requesting streaming token output via CoreML prediction
let config = MLModelConfiguration()
config.computeUnits = .all

let model = try MLModel(contentsOf: modelURL, configuration: config)

// Use MLArrayBatchProvider for token-by-token streaming
var inputTokens: [Int32] = tokenize(prompt)
while !stopConditionMet {
    let inputFeatures = try MLDictionaryFeatureProvider(dictionary: [
        "input_ids": MLMultiArray(inputTokens),
        "attention_mask": currentMask,
        "past_key_values": kvCache
    ])
    let output = try model.prediction(from: inputFeatures)
    let nextToken = argmax(output.featureValue(for: "logits"))
    kvCache = output.featureValue(for: "present_key_values")
    inputTokens = [nextToken]
    emit(nextToken)
}

A 3B parameter model palettized to 4-bit lands around 1.7GB — under the threshold, with headroom left for activations and the KV-cache.

---

The KV-cache growth formula and jetsam

KV-cache memory scales linearly with sequence length:

kv_bytes = 2 × n_layers × n_heads × head_dim × seq_len × dtype_bytes

For a 32-layer, 32-head model with head_dim=128 in float16 at 2048 tokens:

2 × 32 × 32 × 128 × 2048 × 2 = ~1.07 GB

Combined with model weights at 1.7GB, you’re at ~2.8GB total — well past the jetsam threshold on most iPhone configurations. The practical ceiling on iPhone 16 (6GB LPDDR5) is roughly 1024–1500 tokens of usable context before the OS starts issuing memory pressure notifications.

The fix is static KV-cache allocation: pre-allocate the cache tensor at max sequence length and use a position index to avoid dynamic reallocation, which trips memory fragmentation warnings.

---

What to take away

Three things, and you have to get all of them right.

Use ComputeUnit.ALL and profile with Instruments’ Core ML template before overriding anything. Manual ANE-forcing on attention layers reliably hurts latency — the compiler’s scheduling is better than it looks.

Default to 4-bit palettization for weights and float16 for the first and last layers. That combination keeps total weight memory under 2GB while preserving output quality where it matters most: at the model boundaries.

Calculate your KV-cache ceiling before shipping. Use the formula above with your model’s actual architecture constants, then set a hard max_sequence_length in your ML Program spec. Crashing with jetsam in the background is a one-star review waiting to happen.

---

#ios #mobile #architecture #api