---
title: "Wiring CoreML's MLState for Stateful LLM Inference on iPhone"
published: true
description: "Learn how to wire CoreML's MLState API for streaming LLM inference on iPhone — covering KV-cache management, memory pressure eviction, and quantization trade-offs."
tags: swift, ios, mobile, architecture
canonical_url: https://mvpfactory.co/blog/coreml-stateful-llm-inference-iphone
---
## What We Are Building
By the end of this tutorial you will have a working streaming inference pipeline for on-device LLMs using CoreML's `MLState` API. We will wire stateful KV-cache management, handle memory pressure with graceful eviction, enforce a token budget, and gate quantization tier on chip capability at runtime — not on device model strings.
Here is the pattern I use in every on-device LLM project. Miss the KV-cache step and your process will jetsam-kill itself in production after a few minutes of inference. You will not see it in your test suite.
## Prerequisites
- Xcode 16+, targeting iOS 18+
- A CoreML model exported with stateful KV-cache support (`.mlpackage`)
- Familiarity with Swift concurrency (`async/await`, actors, `AsyncThrowingStream`)
- iPhone 14 or later for testing (A15 Bionic minimum)
## Step 1 — Wire MLState from the First Prediction Call
Without state persistence, every token generation step re-feeds the entire context window. Compute cost is quadratic in context length, and KV tensor allocations grow with every prompt until memory pressure terminates your process mid-session. `MLState` attaches mutable state buffers directly to the model, persisted across `prediction(from:using:)` calls.
Here is the minimal setup to get this working:
swift
let model = try MyLLM(configuration: MLModelConfiguration())
let state = model.makeState()
let inputFeatures = MyLLMInput(tokens: promptTokens)
let prefillOutput = try model.prediction(input: inputFeatures, using: state)
for _ in 0..<maxNewTokens {
let decodeInput = MyLLMInput(tokens: [lastToken])
let output = try model.prediction(input: decodeInput, using: state)
// output.logits → sample next token
}
The `state` object holds your cache across the decode loop. No manual tensor serialization. No context re-injection.
## Step 2 — Build the Streaming Inference Actor
Wrap the prefill and decode loop in a Swift actor. This serializes `MLState` access while `AsyncThrowingStream` pushes tokens to the UI layer incrementally.
swift
actor InferenceEngine {
private let model: MyLLM
private var state: MLState?
func generate(prompt: [Int]) -> AsyncThrowingStream<String, Error> {
AsyncThrowingStream { continuation in
Task {
do {
self.state = model.makeState()
let prefillInput = MyLLMInput(tokens: prompt)
_ = try model.prediction(input: prefillInput, using: self.state!)
var lastToken = sampleFromLogits(/* prefill output */)
while lastToken != eosTokenId {
let decodeInput = MyLLMInput(tokens: [lastToken])
let output = try model.prediction(input: decodeInput, using: self.state!)
lastToken = sampleFromLogits(output.logits)
continuation.yield(detokenize(lastToken))
}
continuation.finish()
} catch {
continuation.finish(throwing: error)
}
}
}
}
}
## Step 3 — Handle Memory Pressure Before You Write Another Line
Register for `didReceiveMemoryWarningNotification` immediately. Jetsam — iOS's memory reclamation daemon — terminates your process without warning, without `SIGTERM`, without a grace period. It almost never triggers at model load. It triggers after minutes of inference once KV-cache growth crosses the per-process ceiling.
swift
NotificationCenter.default.addObserver(
forName: UIApplication.didReceiveMemoryWarningNotification,
object: nil,
queue: .main
) { [weak self] _ in
self?.evictKVCache()
}
func evictKVCache() {
state = model.makeState()
contextTokenCount = 0
delegate?.didResetContext()
}
Tell the user what happened. A "memory limit reached — context was trimmed" message is recoverable UX. A jetsam kill is not.
## Step 4 — Enforce a Token Budget
Full cache eviction is a last resort. Track the token count and truncate before you approach the jetsam threshold:
swift
let tokenBudget = 1024 // conservative ceiling for A15/A16
if contextTokenCount + newTokens.count > tokenBudget {
let trimmed = contextBuffer.suffix(tokenBudget / 2)
state = model.makeState()
contextTokenCount = 0
try prefill(tokens: Array(trimmed))
}
Gate the budget and quantization tier on chip capability at runtime:
swift
let isHighMemoryDevice = ProcessInfo.processInfo.physicalMemory >= 8 * 1024 * 1024 * 1024
let quantization: QuantizationMode = isHighMemoryDevice ? .int8 : .int4
let tokenBudget: Int = isHighMemoryDevice ? 2048 : 1024
The docs do not mention this, but checking physical memory is more reliable than parsing device model strings across generations.
## Chip and Memory Ceiling Reference
| Device | Chip | RAM | 4-bit max | 8-bit max |
|---|---|---|---|---|
| iPhone 14 / 14 Pro | A15 / A16 Bionic | 6 GB | ~2.5B params | ~1.2B params |
| iPhone 15 | A16 Bionic | 6 GB | ~2.5B params | ~1.2B params |
| iPhone 15 Pro / 16 Pro | A17 Pro / A18 Pro | 8 GB | ~3.5B params | ~1.8B params |
*Estimates account for ~2 GB OS overhead, KV-cache growth at 2K context (~200–400 MB), and tokenizer buffers.*
Ship 4-bit quantized models for broad compatibility (iPhone 14+). Gate 8-bit behind the memory check above.
## Gotchas
**Jetsam fails silently and late.** It almost never appears in your local test run. It surfaces in crash logs after minutes of sustained inference in production. Profile with Instruments → Memory under sustained load, not a single generation pass.
**`MLState` is not thread-safe.** Two concurrent `prediction(input:using:)` calls on the same state object will corrupt your cache. The actor pattern in Step 2 is not optional.
**4-bit vs 8-bit is a memory ceiling question, not a speed question.** For 4-bit, the Neural Engine pipeline is well-optimized across A15, A16, and A17 Pro. For 8-bit, the wider memory bus on A17 Pro and A18 reduces the decode bottleneck — but the increased RAM headroom is the larger practical win.
**Context truncation surprises users.** Add a visible indicator when the context window is trimmed. Invisible resets erode trust fast.
## Conclusion
Stateless inference at generation time is architecturally broken on mobile — compute cost grows with context length, and no battery optimization compensates for it. Wire `MLState` from the start, register your memory warning handler before your first decode loop, and enforce a token budget tuned to the chip generation you are targeting. Everything else in your pipeline follows from those three constraints.
**Further reading:** [CoreML MLState documentation](https://developer.apple.com/documentation/coreml/mlstate), [WWDC24 — Bring your machine learning and AI models to Apple silicon](https://developer.apple.com/videos/play/wwdc2024/10159/)
Top comments (0)