CoreML stateful LLM inference on iPhone: KV-cache & memory
SEO Meta 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.
TL;DR
CoreML’s MLState API (iOS 18+) enables persistent KV-cache across prediction calls, unlocking true stateful autoregressive inference on-device. The hard constraint is jetsam: exceed the memory ceiling and your process dies silently. 4-bit quantization is the only practical path to models above 1B parameters on iPhone 14 hardware. A17 Pro and A18 unlock 8-bit at the same parameter counts due to increased RAM headroom. Design your pipeline to evict cache gracefully on didReceiveMemoryWarning — or your UX will crater in production.
Why stateful inference matters
A 2B-parameter model that performs well on the first 10 prompts can jetsam-kill your process by the 11th.
Without state persistence, every token generation step re-feeds the entire context window. On server-side, this is wasteful. On iPhone, it’s a dealbreaker: compute cost is quadratic in context length, and KV tensor allocations grow with every prompt until memory pressure terminates your process mid-session.
The KV-cache stores pre-computed key and value tensors for all prior tokens. CoreML’s MLState API solves the re-computation problem by attaching mutable state buffers directly to the model, persisted across prediction(from:using:) calls.
let model = try MyLLM(configuration: MLModelConfiguration())
let state = model.makeState()
// Prefill pass — encode the prompt
let inputFeatures = MyLLMInput(tokens: promptTokens)
let prefillOutput = try model.prediction(input: inputFeatures, using: state)
// Decode loop — state carries the KV-cache forward
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. No manual tensor serialization. No context re-injection. This is the primitive you build your streaming pipeline on.
Architecture: the streaming inference pipeline
Let me walk you through the architecture that holds up in production.
TokenQueue → Tokenizer → Prefill Engine
↓
[MLState / KV-Cache]
↓
Decode Loop (token-by-token)
↓
StreamingTextDelegate → UITokenQueue is a serial queue that serializes incoming prompt requests — worth the overhead when your UI allows mid-generation interruption or prompt chaining. It ensures only one prefill runs at a time against a shared MLState instance.
The prefill and decode passes run on separate DispatchQueue lanes with a shared actor protecting MLState access. Swift concurrency (async/await + actors) maps cleanly here: the state actor serializes access while async streams push tokens to the UI layer incrementally.
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)
}
}
}
}
}Errors propagate through the AsyncThrowingStream — the call site handles CoreML prediction failures and memory-related throws without crashing the actor.
Memory pressure: jetsam is not negotiable
Jetsam, iOS’s memory reclamation daemon, will terminate your process without warning if you cross the per-process memory limit. No SIGTERM, no grace period — just a silent kill and a crash log. The deceptive failure mode: this almost never happens at model load. It happens after minutes of inference, once KV-cache growth pushes you over the threshold. By then you’re reading crash reports, not test logs.
Chip and RAM comparison
| Device | Chip | RAM | 4-bit max model | 8-bit max model |
|---|---|---|---|---|
| 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 OS overhead (~2 GB reserved), model weights, KV-cache growth at 2K context (~200–400 MB depending on head count), and tokenizer buffers. A17 Pro and A18 Pro rows are collapsed — the 8-bit ceiling difference between them falls within estimation noise and should not be treated as a meaningful distinction without device-specific profiling.
Graceful cache eviction
Register for memory warnings and evict aggressively:
NotificationCenter.default.addObserver(
forName: UIApplication.didReceiveMemoryWarningNotification,
object: nil,
queue: .main
) { [weak self] _ in
self?.evictKVCache()
}
func evictKVCache() {
// Reset state — user will experience a context reset,
// but the process survives.
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.
Context window and token budget management
Evicting the full KV-cache is a last resort. A more controlled strategy is enforcing a rolling token budget: track the token count in the active context and truncate before you approach the jetsam threshold.
let tokenBudget = 1024 // conservative ceiling for A15/A16
if contextTokenCount + newTokens.count > tokenBudget {
// Truncate oldest tokens, reset state, re-prefill from truncated context
let trimmed = contextBuffer.suffix(tokenBudget / 2)
state = model.makeState()
contextTokenCount = 0
try prefill(tokens: Array(trimmed))
}This gives the user a continuous experience while keeping KV-cache size bounded. The truncation point is tunable per chip generation — tighter on A15/A16, more generous on A17 Pro and A18 hardware.
Quantization trade-offs: 4-bit vs. 8-bit on Neural Engine
The A17 Pro and A18 introduced meaningful Neural Engine throughput improvements. For 4-bit models, the Neural Engine pipeline is well-optimized across all three chip generations. For 8-bit, the wider memory bus on A17 Pro and A18 reduces the decode bottleneck — but the memory ceiling improvement is the larger practical win, not raw tokens-per-second.
The practical recommendation: ship 4-bit quantized models for broad compatibility (iPhone 14+), and offer an 8-bit quality tier gated behind a runtime chip check.
Putting it together
Stateless inference at generation time is architecturally broken on mobile — compute cost grows with context length, and no battery optimization compensates. Wire MLState from the start.
Everything else follows from the memory constraint. Wire didReceiveMemoryWarningNotification to a cache reset path before you write your first decode loop. Enforce a token budget so normal usage never touches that limit.
Gate quantization and token budget on chip at runtime, not device model strings:
func neuralEngineGeneration() -> Int {
var size = 0
sysctlbyname("hw.optional.arm.FEAT_BF16", nil, &size, nil, 0)
// A17 Pro+ supports BF16 natively; use as a capability proxy
var supported: Int32 = 0
sysctlbyname("hw.optional.arm.FEAT_BF16", &supported, &size, nil, 0)
return Int(supported)
}
// Or via ProcessInfo for a simpler heuristic:
let isHighMemoryDevice = ProcessInfo.processInfo.physicalMemory >= 8 * 1024 * 1024 * 1024
let quantization: QuantizationMode = isHighMemoryDevice ? .int8 : .int4
let tokenBudget: Int = isHighMemoryDevice ? 2048 : 1024Serve 4-bit with a 1K token budget to A15/A16 hardware, 8-bit with a 2K budget to A17 Pro and above. This keeps your memory ceiling conservative and avoids the late-session jetsam failure that only surfaces after minutes of inference — the one that ships to production and shows up in crash logs, not in your test suite.
Tags: ios swift mobile architecture api