Streaming Core ML LLM tokens to SwiftUI with actors

Meta description: Stream Core ML LLM tokens to SwiftUI token-by-token using Swift 6 actors, AsyncStream back-pressure, and ANE dispatch tuning to sustain 15+ tok/s on device.

---

TL;DR

Running a local LLM on-device is not the hard part. Streaming its output token-by-token into a responsive SwiftUI view without blocking the main thread, blowing the memory budget, or stalling on prefill — that’s where most implementations fall apart. This post walks through the full pipeline: Core ML → Swift actor → AsyncStream → SwiftUI, with the specific tradeoffs that determine whether you hit 15 tok/s or a memory termination.

Tags: ios swift swiftui mobile architecture

---

The architecture at a glance

The pipeline has four layers, each with a strict ownership contract:

CoreMLEngine (Actor)
    └── AsyncStream<Token>
            └── TokenStreamViewModel (@MainActor)
                    └── SwiftUI Text view

The actor owns inference state, the stream applies back-pressure, the view model bridges onto the main actor. Violate any boundary and you get either a data race under Swift 6’s strict concurrency checking, or a dropped frame.

---

The actor: owning inference state

actor CoreMLEngine {
    private let model: MLModel
    // Pre-allocated once; mutated in place every decode step
    private var inputArray: MLMultiArray
    private var kvCache: MLMultiArray

    init(model: MLModel) throws {
        self.model = model
        self.inputArray = try MLMultiArray(shape: [1, 1], dataType: .int32)
        self.kvCache    = try MLMultiArray(shape: kvCacheShape, dataType: .float16)
    }

    func generateTokens(prompt: String) -> AsyncStream<String> {
        AsyncStream(bufferingPolicy: .bufferingNewest(16)) { continuation in
            Task {
                var tokenIds = tokenize(prompt)
                while !shouldStop(tokenIds) {
                    guard let next = runSingleStep(context: tokenIds) else { break }
                    tokenIds.append(next.id)
                    continuation.yield(next.text)
                }
                continuation.finish()
            }
        }
    }

    // One autoregressive step: feed the latest token, return the next one
    private func runSingleStep(context: [Int]) -> (id: Int, text: String)? {
        inputArray[0] = context.last.map(NSNumber.init) ?? 0
        let features = try? MLDictionaryFeatureProvider(dictionary: [
            "input_ids":       MLFeatureValue(multiArray: inputArray),
            "past_key_values": MLFeatureValue(multiArray: kvCache)
        ])
        guard let features,
              let prediction   = try? model.prediction(from: features),
              let logits       = prediction.featureValue(for: "logits")?.multiArrayValue,
              let updatedCache = prediction.featureValue(for: "present_key_values")?.multiArrayValue
        else { return nil }
        kvCache = updatedCache   // discard old buffer, take ownership of model output
        let id = argmax(logits)
        return (id: id, text: detokenize(id))
    }
}

The actor boundary is load-bearing in Swift 6. MLModel, inputArray, and kvCache all live here and nowhere else, cutting out the data races that plagued pre-concurrency Core ML integrations.

---

MLMultiArray pre-allocation and KV cache threading

Most teams miss this: every decode step allocates a new input MLMultiArray unless you deliberately reuse one. Across a 512-token generation run, that is 512 short-lived heap objects, measurable allocator pressure that compounds under thermal throttling.

The deeper problem is unbounded KV cache growth. If your compiled model exposes past_key_values as an input/output pair (the standard pattern for autoregressive Core ML exports), you must pass the cache back in on every step. Failing to do so forces the model to recompute the full attention history from scratch each token: prefill cost at decode time, with guaranteed OOM on sequences beyond roughly 128 tokens.

// Wrong: reallocates every step, forces full-context attention recompute
private func runSingleStep_bad() {
    let input = try! MLMultiArray(shape: [1, 1], dataType: .int32)
    // no past_key_values — model recomputes full context on every call
}

// Correct: pre-allocated input, explicit cache threading (see actor init above)

In my experience building production systems, missing KV cache threading is the single most common reason an on-device LLM that passes a 50-token smoke test terminates with a jetsam kill at 300 tokens.

---

ANE vs. CPU: the dispatch tradeoff

Dispatch Target	Throughput (tok/s)	Prefill Latency	Memory Footprint
ANE only	15-22 tok/s	180ms	Low (weights cached)
CPU only	4-7 tok/s	90ms	High (no dedicated cache)
ANE + CPU split	12-18 tok/s	120ms	Medium

ANE wins on throughput for autoregressive decode. CPU wins on prefill latency for short prompts. The practical strategy: compile with computeUnits = .cpuAndNeuralEngine and let the runtime decide per-operation. For models above 2B parameters, force .neuralEngine for attention layers explicitly via model metadata, since the automatic scheduler makes conservative choices under memory pressure.

To handle memory warnings with correct Swift 6 actor isolation:

nonisolated func observeMemoryPressure() {
    NotificationCenter.default.addObserver(
        forName: UIApplication.didReceiveMemoryWarningNotification,
        object: nil,
        queue: nil
    ) { [weak self] _ in
        guard let self else { return }
        Task { await self.flushKVCache() }
    }
}

The method is nonisolated so it can be called from any context without crossing an actor boundary. The closure captures self weakly (actors are reference types) and re-enters actor isolation via Task { await } rather than optional-chaining an actor-isolated method from a non-isolated context, which Swift 6’s strict concurrency checker correctly rejects.

---

Back-pressure with AsyncStream

AsyncStream does not apply back-pressure by default. If the consumer cannot keep up, tokens buffer unboundedly in the continuation’s internal queue. On an A17 Pro, the ANE can outrun SwiftUI’s text layout engine by 3-4x during burst decode.

Setting a bufferingPolicy caps that growth:

AsyncStream(String.self, bufferingPolicy: .bufferingNewest(16)) { continuation in
    // inference loop
}

Sixteen tokens is sufficient headroom for rendering jitter without letting the queue grow proportional to model throughput.

---

SwiftUI without main-thread blocking

@MainActor
class TokenStreamViewModel: ObservableObject {
    @Published var output = ""

    func stream(from engine: CoreMLEngine, prompt: String) async {
        for await token in await engine.generateTokens(prompt: prompt) {
            output += token
        }
    }
}

@MainActor on the view model combined with for await on the stream gives cooperative scheduling. Each await is a suspension point — SwiftUI’s render loop gets CPU time between tokens. No DispatchQueue.main.async required.

---

The memory budget that fits on device

For a 1B-parameter INT4 quantized model: expect ~600MB weights, ~120MB KV cache at 512 context, ~40MB runtime overhead — ~760MB peak, under the ~1.2GB soft limit on an iPhone 15 before jetsam intervenes.

Beyond 2B INT4, you are in pressure territory on older devices. Profile with Instruments’ Memory Graph, not Xcode’s summary, which underreports ANE allocations.

---

Three things worth knowing

1. Pre-allocate MLMultiArray for inputs and thread your KV cache explicitly on every decode step. Per-step reallocation creates heap churn; a missing past_key_values input forces full-context attention recomputation and guarantees OOM on sequences above ~128 tokens.

2. Set bufferingPolicy: .bufferingNewest(N) on your AsyncStream. Without it you have an unbounded queue that masks latency problems in development and surfaces them under real thermal conditions in production.

3. Profile ANE dispatch with Instruments’ Core ML template before shipping. The default computeUnits selection is conservative — manually targeting .neuralEngine for attention layers can recover 30-40% throughput on supported hardware.