Real-time diffusion on iOS: CoreML, KV-cache & interactive editing

Meta description: Ship quantized Stable Diffusion on iOS with CoreML ML Program format, cross-attention KV-cache reuse, and per-chip memory-tier fallback for A16, A17, and M-series.

---

TL;DR

Shipping a real-time diffusion-based image editor on iOS is achievable — but only if you treat memory pressure as a first-class constraint, not an afterthought. The UI canvas runs at 60fps because inference executes asynchronously off the main thread; actual per-step latency runs from 280ms to 800ms depending on chip tier and quantization depth. The difference between a smooth interactive experience and a jittery, thermally-throttled one comes down to three decisions: model quantization depth, attention KV-cache reuse across denoising steps, and a hard per-chip memory ceiling that triggers quality-tier fallback before the OS kills your process.

---

The architecture challenge

What most teams get wrong about on-device diffusion: they benchmark on a plugged-in M2 iPad and ship to A16 iPhones. The Neural Engine tier gap is brutal — not just in raw TOPS, but in the amount of on-chip SRAM available for intermediate activations before the runtime spills to main memory and your latency doubles.

Latent diffusion on iOS breaks into four compilable segments: the text encoder, a VAE encoder, the U-Net denoising backbone, and a VAE decoder. Each is compiled separately into CoreML’s .mlpackage (ML Program format). This separation isn’t optional ceremony — it’s the only way to cache and reuse the U-Net’s cross-attention key/value tensors across denoising steps without recomputing them from scratch on every iteration.

---

Quantization: picking your precision floor

A full FP32 SD 1.5 U-Net is unusable on-device. FP16 is the starting point; palettized INT8 and INT4 weight compression via coremltools.optimize.coreml.palettize_weights is where it gets interesting.

Precision	U-Net Size	ANE Eligible	Typical Step Latency (A17)
FP16	~2.5 GB	Partial	~800 ms/step
INT8 (weights)	~1.3 GB	Yes	~420 ms/step
INT4 (weights)	~700 MB	Yes	~280 ms/step
INT4 + attention FP16	~750 MB	Yes	~295 ms/step

The last row is the production choice. Keep attention projections at FP16 — palettizing them to INT4 compounds error across 20 denoising steps in ways that are visually obvious. Palettize everything else.

---

Attention KV-cache reuse across denoising steps

This is the highest-leverage optimization most iOS ML engineers skip. In a guided diffusion edit (inpainting, style transfer), the text conditioning doesn’t change between steps. That means the cross-attention K and V projections from your text encoder output are identical on every step.

// Compile U-Net with stateful KV-cache via MLProgram
let config = MLModelConfiguration()
// computeUnits is set per chip tier — see next section
config.computeUnits = resolvedComputeUnits()

// Cache K/V outputs as MLMultiArray across steps
var cachedKV: [String: MLMultiArray] = [:]

func denoisingStep(latent: MLMultiArray, step: Int) throws -> MLMultiArray {
    var inputDict: [String: Any] = [
        "latent_input": latent,
        "timestep": MLMultiArray([step]),
        "use_cached_kv": MLMultiArray([step > 0 ? 1 : 0])
    ]

    // Inject cached K/V tensors on steps > 0
    if step > 0 {
        for (key, value) in cachedKV { inputDict[key] = value }
    }

    do {
        let provider = try MLDictionaryFeatureProvider(dictionary: inputDict)
        let output = try unet.prediction(from: provider)

        if step == 0 {
            cachedKV = extractKV(from: output) // returns [String: MLMultiArray]
        }

        guard let result = output.featureValue(for: "latent_output")?.multiArrayValue else {
            throw InferenceError.missingOutput
        }
        return result
    } catch {
        throw error
    }
}

In my experience building production inference pipelines, this alone cuts cross-attention compute by 35–45% on a 20-step schedule — enough to meaningfully reduce total generation time and keep the UI responsive during async inference.

---

The memory pressure ceiling by chip tier

This is where production deployments fail silently. iOS won’t crash your app immediately when you breach the Neural Engine’s working-set limit — it will silently delegate layers to CPU, which is 4–8x slower and shows up as dropped frames, not errors.

Chip	ANE Working Set Budget	Safe Model Budget	Fallback Trigger
A16 Bionic	~1.0 GB	~700 MB	CPU delegation above ~1.1 GB
A17 Pro	~1.4 GB	~1.0 GB	CPU delegation above ~1.5 GB
M2 / M4 (iPad)	~3.5 GB	~2.5 GB	Rarely triggered

The MLModelConfiguration.computeUnits setting should reflect the chip tier detected at runtime — don’t hardcode .cpuAndNeuralEngine universally:

func resolvedComputeUnits() -> MLComputeUnits {
    let chip = ChipTierDetector.current() // wrapper around sysctlbyname("hw.optional.*")
    switch chip {
    case .a16:
        // Constrained ANE budget — avoid CPU+GPU+ANE contention
        return .cpuAndNeuralEngine
    case .a17:
        return .cpuAndNeuralEngine
    case .m2, .m4:
        // Larger unified memory pool; `.all` enables GPU path for non-ANE ops
        return .all
    default:
        return .cpuAndNeuralEngine
    }
}

Monitor os_proc_available_memory() before each inference pass. If headroom drops below your model’s activation footprint, drop to a lower-resolution latent space (e.g., 384×384 instead of 512×512) rather than letting the runtime silently degrade.

---

Three takeaways

1. Palettize weights to INT4, keep attention at FP16. This is the precision configuration that fits within A16/A17 ANE budgets while maintaining generation quality across 20+ denoising steps. Use coremltools.optimize.coreml.palettize_weights with a 4-bit config and exempt attention projection layers explicitly.

2. Cache cross-attention K/V tensors on step zero. For any conditioning that doesn’t change between steps, recomputing K/V projections is pure waste — structure your ML Program to accept cached inputs and gate computation on a timestep flag. That’s a 35–45% compute reduction with no quality cost.

3. Instrument memory headroom and set computeUnits per chip tier. Silent CPU delegation is your real enemy. Detect chip tier at runtime, configure MLModelConfiguration accordingly, check available memory before inference, and fail gracefully to lower resolution rather than letting the ANE runtime make that decision for you.

---

Tags: ios, swift, mobile, architecture, coreml