Wiring iOS CoreML's Stateful Models to a Streaming Inference Pipeline

The guide explains how to set up a streaming inference pipeline for on‑device large language models with CoreML’s MLState API. It covers KV‑cache wiring, memory‑warning handling, token‑budget enforcement, and quantization decisions based on device memory. The article also lists common pitfalls and…

Tue, 15 Sep 2026 07:52:14 GMT

When you want to run a large language model (LLM) directly on an iPhone, the naive approach of feeding the entire prompt to the model for every token is both slow and memory‑intensive. Apple’s MLState API solves this by keeping a mutable state buffer inside the model, so that each new token can be generated from the previous one without re‑allocating the entire key‑value cache. This article walks through the steps you need to build a robust, streaming inference pipeline that works on iPhone 14 and newer, handles memory pressure gracefully, and picks the right quantization level for each device.

What We’re Building

By the end of this tutorial you’ll have a fully functional streaming inference loop that:

• Initializes MLState on the first prediction call.
• Maintains a persistent KV cache across all subsequent calls.
• Responds to iOS memory‑warning notifications by evicting the cache.
• Keeps the token count below a device‑specific budget to avoid jetsam termination.
• Selects 4‑bit or 8‑bit quantization at runtime based on available RAM.

Prerequisites

You’ll need:

• Xcode 16 or later, targeting iOS 18+.
• A CoreML model exported with stateful KV‑cache support (the .mlpackage format).
• Basic familiarity with Swift concurrency (async/await, actors, AsyncThrowingStream).
• An iPhone 14 or newer for testing (A15 Bionic is the minimum required chip).

Step 1 – Wire MLState from the First Prediction Call

Without state persistence, every token generation step would re‑feed the entire context window, resulting in quadratic compute cost and continuous KV tensor growth. MLState attaches a mutable buffer to the model that survives across prediction(from:using:) calls. Below is a minimal example that shows how to create the state, run a pre‑fill, and then stream tokens until the end‑of‑sentence token is produced.

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... {
    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(prefillOutput.logits)
                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 2 – Handle Memory Pressure Early

iOS’s jetsam daemon can terminate your app without warning once the process exceeds its memory ceiling. It rarely fires during model load but can kill the app after a few minutes of inference when the KV cache grows too large. Register for UIApplication.didReceiveMemoryWarningNotification immediately after creating the model and reset the state when the notification fires.

NotificationCenter.default.addObserver(
    forName: UIApplication.didReceiveMemoryWarningNotification,
    object: nil,
    queue: .main
) { [weak self] _ in
    self?.evictKVCache()
}

func evictKVCache() {
    state = model.makeState()
    contextTokenCount = 0
    delegate?.didResetContext()
}

When you evict the cache, inform the user with a message such as “Memory limit reached – context was trimmed.” This gives the user a clear indication that the conversation was reset, rather than a silent crash.

Step 3 – Enforce a Token Budget

Full cache eviction should be a last resort. Instead, track the number of tokens in the current context and trim it before you hit the jetsam threshold. The budget varies by chip generation, so you can determine it at runtime by inspecting the device’s physical memory.

let isHighMemoryDevice = ProcessInfo.processInfo.physicalMemory >= 8 * 1024 * 1024 * 1024
let quantization: QuantizationMode = isHighMemoryDevice ? .int8 : .int4
let tokenBudget: Int = isHighMemoryDevice ? 2048 : 1024

if contextTokenCount + newTokens.count > tokenBudget {
    let trimmed = contextBuffer.suffix(tokenBudget / 2)
    state = model.makeState()
    contextTokenCount = 0
    try prefill(tokens: Array(trimmed))
}

The table below summarizes the maximum model size you can run on each recent iPhone model without triggering jetsam, assuming 2‑bit overhead for the OS and tokenizer buffers.

Ship 4‑bit quantized models for broad compatibility (iPhone 14+). Gate 8‑bit models behind the memory check above to avoid jetsam on lower‑end devices.

Common Gotchas

• Jetsam kills silently; it rarely appears in local tests but shows up in crash logs after sustained inference.
• MLState is not thread‑safe. Run prediction calls sequentially or wrap them in an actor.
• 4‑bit vs 8‑bit is a memory ceiling question, not a speed question. 8‑bit models benefit from the wider memory bus on newer chips.
• Context truncation can surprise users. Always display a visible indicator when the context window is trimmed.

Conclusion

Stateless inference on mobile is fundamentally inefficient because compute cost scales with context length. By wiring MLState from the start, registering a memory‑warning handler before the first decode loop, and enforcing a token budget tuned to the target chip, you can build a reliable, on‑device LLM experience that stays within the device’s memory limits.

Further Reading

• CoreML MLState documentation
• WWDC24 – Bring your machine learning and AI models to Apple silicon

Why it matters

On‑device LLMs offer privacy and low‑latency interactions, but they require careful memory management to avoid crashes. This guide shows how to keep your app stable while delivering real‑time language generation on the latest iPhones.

Key points

• Use MLState to persist KV cache across predictions
• Register for memory‑warning notifications before inference starts
• Trim context before hitting jetsam by enforcing a token budget
• Choose 4‑bit or 8‑bit quantization based on available RAM
• Avoid concurrent predictions on the same state object

Frequently asked questions

What is jetsam and why does it kill my app?

Jetsam is iOS’s memory reclamation daemon that terminates processes that exceed their memory ceiling. It does so without warning, so apps that grow their KV cache too large during inference can be killed after a few minutes of use.

How do I know which quantization level to use on a device?

Check the device’s physical memory at runtime. Devices with 8 GB of RAM can safely run 8‑bit models; those with 6 GB should use 4‑bit to stay below the jetsam threshold.

Is MLState thread‑safe?

No. If you need concurrent predictions, wrap the state in an actor or serialize calls to avoid cache corruption.

Reporting drawn from

• dev.to