---
title: "Flash Attention on Mobile: Wiring CoreML's Multi-Head Attention to ANE for Sub-20ms Prefill on Long Contexts"
published: true
description: "Learn how to implement tiled attention on Apple Neural Engine using CoreML stateful primitives for fast prefill on long-context prompts for on-device LLMs."
tags: ios, mobile, architecture, swift
canonical_url: https://mvpfactory.co/blog/flash-attention-coreml-ane-prefill
---
## What We Will Build
By the end of this workshop, you will understand why naive attention collapses on Apple Neural Engine beyond 2K tokens, and you will have a working chunked prefill loop using CoreML's stateful `MLState` primitives. This is the pattern I use in every on-device LLM integration — not a toy demo, but the structure you need for production inference pipelines.
**Prerequisites:**
- iOS 18+ / macOS 15+ deployment target
- Xcode 16+
- A CoreML-converted LLM (4-bit or 8-bit quantized)
- Basic familiarity with `MLModel` and `MLMultiArray`
---
## The Problem: Naive Attention Destroys ANE Performance
Most teams benchmark single-token decode latency, ship something that looks fast, and then discover their prefill on a 1,500-token system prompt takes seconds. The user experience collapses on the first real-world session.
Here is the math that explains why. Standard multi-head attention over a sequence of length *N* materializes an *N × N* score matrix. At 2K tokens with 32 heads in float16, that is north of 250MB of intermediate activation. The ANE cannot hold that in its local memory hierarchy — it forces expensive round-trips to DRAM.
ANE-to-DRAM bandwidth is constrained relative to the ANE's compute throughput. When attention spills to DRAM, you lose the hardware advantage you were optimizing for. High raw throughput becomes irrelevant.
---
## The Flash Attention Insight
Flash Attention (Dao et al., 2022) reorders attention computation so the full *N × N* matrix is never materialized. You tile query, key, and value tensors into blocks, compute softmax incrementally using online normalization, and accumulate results — all within fast on-chip memory.
for each tile of Q:
for each tile of K, V:
compute local attention scores
update running max and sum for online softmax
accumulate weighted V into output tile
On GPU this maps to shared memory. On ANE, the analog is keeping tile intermediates within the ANE's neural memory — avoiding DRAM round-trips for the attention matrix itself. The algorithm trades recomputation for memory efficiency, and on ANE hardware that trade is almost always worth it.
---
## Step-by-Step: CoreML Stateful Prefill
CoreML's stateful model support gives you explicit KV-cache management via `MLState`. This is what makes chunked prefill practical on ANE without custom Metal kernels.
> **The docs do not mention this clearly, but** `MLState` and stateful CoreML models require **iOS 18+ / macOS 15+ and Xcode 16+**. Confirm your deployment target before adopting this pattern — it will silently fall back otherwise.
Here is the minimal setup to get this working:
swift
// Configure stateful KV-cache model
let config = MLModelConfiguration()
config.computeUnits = .cpuAndNeuralEngine
let model = try MLModel(contentsOf: modelURL, configuration: config)
let kvState = try model.makeState()
// Chunked prefill loop
let chunkSize = 256 // tuned to ANE working memory
for chunkStart in stride(from: 0, to: promptTokens.count, by: chunkSize) {
let chunkSlice = Array(promptTokens[chunkStart..<min(chunkStart + chunkSize, promptTokens.count)])
let seqLen = chunkSlice.count
// MLMultiArray requires explicit shape and dataType
let tokenArray = try MLMultiArray(shape: [1, NSNumber(value: seqLen)], dataType: .int32)
let posArray = try MLMultiArray(shape: [1, NSNumber(value: seqLen)], dataType: .int32)
for i in 0..<seqLen {
tokenArray[i] = NSNumber(value: chunkSlice[i])
posArray[i] = NSNumber(value: chunkStart + i)
}
let input = try MLDictionaryFeatureProvider(dictionary: [
"input_ids": tokenArray,
"position_ids": posArray
])
let _ = try model.prediction(from: input, using: kvState)
}
The `MLState` object persists KV entries across chunk calls. Each chunk adds to the cache without re-processing prior tokens — exactly the structure Flash Attention's tiling requires.
---
## Tuning Chunk Size
Chunk size is a tuning exercise, not a fixed answer. Here is the gotcha that will save you hours: too small and dispatch overhead dominates; too large and you overflow ANE memory and fall back to DRAM-bound execution.
| Chunk Size (tokens) | ANE Utilization | DRAM Pressure | Prefill Throughput |
|---|---|---|---|
| 64 | Low (dispatch overhead) | Minimal | Moderate |
| 128 | Moderate | Low | Good |
| 256 | High | Low | **Best (typical)** |
| 512 | High–saturated | Moderate | Varies by model |
| 1024 | May spill | High | Degrades |
256 tokens is a consistent sweet spot across modern LLM architectures quantized to 4-bit or 8-bit weights on iPhone-class hardware. Profile your specific model with Instruments' Core ML template before committing — this table is a starting point, not a guarantee.
---
## Gotchas
**Quantization is a prerequisite, not an afterthought.** No amount of attention tiling compensates for a model that does not fit ANE constraints. CoreML's ANE backend requires weights in a format it can load into neural memory — typically 4-bit or 8-bit palettized or linear quantized, converted via `coremltools` `ct.optimize`. Teams that retrofit quantization late consistently discover ANE fallback to CPU for layers the compiler cannot schedule.
**Profile compute unit assignments before shipping.** Run `xcrun coremlcompiler compile` and inspect the output. If layers are falling back to CPU, tiling strategy is irrelevant — fix the model first.
**Measure prefill separately from decode.** Decode speed is table stakes; prefill over long system prompts is where most shipped apps fall down. Instrument both independently before declaring success.
**`MLState` does not auto-reset between sessions.** If you reuse a state object across user sessions without reinitializing, you will accumulate stale KV entries and corrupt attention outputs. Create a fresh state per inference session.
---
## Conclusion
Let me show you the pattern I use in every on-device LLM project: start at 256 tokens per chunk, use `MLState` for stateful KV-cache accumulation, profile with Core ML Instruments to confirm ANE utilization stays high and DRAM pressure stays low, then adjust. Quantize before you optimize attention — without 4-bit or 8-bit weights via `coremltools`, the model will not schedule to ANE at all, and everything else is irrelevant.
The payoff is real: chunked prefill with stateful KV-cache restructures long-context attention into tiles the ANE can actually handle, bypassing the DRAM bottleneck that tanks naive implementations above 2K tokens.
---
## References
- Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). *FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness*. [arXiv:2205.14135](https://arxiv.org/abs/2205.14135)
- [Apple CoreML Documentation — Stateful Models](https://developer.apple.com/documentation/coreml)
- [coremltools Optimization Guide](https://apple.github.io/coremltools/docs-guides/source/optimization-overview.html)
Top comments (0)