Wiring Apple's Neural Engine to Core ML's Stateful Models

---
title: "Persistent KV-Cache on iOS 18: Wiring Core ML Stateful Models to the Neural Engine"
published: true
description: "Core ML stateful models let you persist KV-cache tensors across inference calls, killing the O(n²) prefill cost that destroys on-device LLM chat performance."
tags: ios, swift, mobile, architecture
canonical_url: https://mvpfactory.co/blog/core-ml-stateful-kv-cache-ios18
---

What We Are Building

By the end of this tutorial you will have an on-device chat inference session that reuses key-value cache tensors across turns instead of re-prefilling the full context each time. We will wire Core ML's stateful model API, pin execution to the ANE, handle conversation resets cleanly, and use a Swift 6 actor to prevent a subtle race condition that passes casual testing and destroys you in production.

---

Prerequisites

• Xcode 16+ targeting iOS 18
• A Core ML .mlpackage ready for conversion (or an existing stateful model)
• coremltools 8+ for the Python conversion step
• Familiarity with MLModel and MLMultiArray

---

Step 1 — Declare KV-Cache Buffers as Stateful at Conversion Time

Here is the pattern I use in every project. The states= parameter is not optional — nothing downstream works without it.

import coremltools as ct

kv_cache_state = ct.StateType(
    wrapped_type=ct.TensorType(
        shape=(num_layers, num_heads, max_seq_len, head_dim)
    ),
    name="kv_cache"
)

mlmodel = ct.convert(traced_model, states=[kv_cache_state], ...)
mlmodel.save("chat_model.mlpackage")

Pass it as a list, not a dict. The shape encodes your max_seq_len — bake this in consciously because you will need it later.

---

Step 2 — Create One MLState and Reuse It

At inference time, create a single MLState and pass it into every prediction call. The cache tensors read and write in-place — no copy, no re-allocation per turn.

let state = try model.makeState()

func generateNextToken(inputIds: MLMultiArray, cachePosition: Int) throws -> MLMultiArray {
    let input = ChatModelInput(input_ids: inputIds, cache_position: cachePosition)
    let output = try model.prediction(from: input, using: state)
    return output.logits
}

---

Step 3 — Force ANE Execution, Exclude the GPU

The stateful read/write path uses scatter-gather memory access patterns the GPU handles poorly. Left unconfigured, Core ML falls back to GPU for state update ops and you lose most of the latency gain.

let config = MLModelConfiguration()
config.computeUnits = .cpuAndNeuralEngine  // Exclude GPU explicitly
let model = try MLModel(contentsOf: modelURL, configuration: config)

Here is what that costs you on iOS 18 / A17 Pro at 512 tokens sustained load:

Compute Units	First-token latency	Per-token latency	Thermal impact
.all (default)	~210ms	~38ms	High
.cpuAndNeuralEngine	~180ms	~22ms	Low
.cpuOnly	~640ms	~95ms	Minimal

That ~40% per-token improvement compounds across a long conversation. Thermal throttling will eat those gains back if you ignore it.

---

Step 4 — Handle Resets and Sequence Overflow

MLState does not reset automatically. When the user starts a new conversation, recreate the state object:

actor InferenceSession {
    private var state: MLState
    private let model: MLModel

    init(model: MLModel) throws {
        self.model = model
        self.state = try model.makeState()
    }

    func resetConversation() throws {
        state = try model.makeState()  // Fresh zero-filled buffers
    }
}

Do not try to zero-fill the buffers manually. The internal layout of MLState is opaque and may change across OS versions.

When cache_position approaches max_seq_len, use the summarize-and-reset pattern — it is the only production-safe strategy:

func generateNextToken(inputIds: MLMultiArray, cachePosition: Int) throws -> MLMultiArray {
    if cachePosition >= maxSeqLen - safetyMargin {
        try await summarizeAndReset()
    }
    // ... normal inference
}

The actor wrapper above also serializes concurrent calls to generateNextToken automatically. Swift 6 strict concurrency will flag the race at compile time if you skip this — do not ignore that warning.

---

Gotchas

states= must be a list, not a dict. The docs do not always make this obvious. Passing a dict silently produces a model that ignores the state at runtime.

MLState survives until deallocation. If you keep a long-lived InferenceSession, stale KV tensors from a previous conversation will corrupt the next one. Always recreate on reset.

GPU fallback is silent. Core ML will not warn you if it routes ops to the GPU. Verify with Instruments → Core ML profiling that kv_cache ops land on the ANE.

The Swift 6 actor boundary is load-bearing. A concurrent write to a cache tensor produces output that looks plausible. You will not catch it in testing. You will catch it when a user reports incoherent responses mid-conversation.

---

Conclusion

Stateful Core ML models give you server-class KV-cache behaviour on-device. The three things that actually matter: declare buffers via states=[...] at conversion, lock compute units to cpuAndNeuralEngine, and wrap your session in a Swift 6 actor. Get those right and you have an on-device chat loop that scales with conversation length instead of collapsing under it.