Flash attention on ANE: fast long-context prefill

Meta 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.

---

TL;DR

Running attention over long contexts on-device is a memory bandwidth problem, not a compute problem. The ANE has high raw throughput, but naive full-sequence attention blows past its working memory limits above ~2K tokens. Chunked KV-cache prefill using CoreML’s stateful attention primitives restructures the problem into tiles the ANE can actually handle — real prefill speedups without touching the GPU.

---

The problem with naive attention on ANE

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.

The ANE is purpose-built for matrix operations with high operational intensity. But standard multi-head attention over a sequence of length N materializes an N × N attention score matrix. At 2K tokens with 32 heads and float16, that’s north of 250MB of intermediate activation — more than the ANE can hold in its local memory hierarchy, forcing 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.

---

Flash attention: the core insight

Flash Attention (Dao et al., 2022) reorders the attention computation so the full N × N matrix is never materialized. Instead, you tile the query, key, and value tensors into blocks, compute softmax incrementally using online normalization, and accumulate results — all within fast on-chip memory.

The algorithm trades recomputation for memory efficiency:

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 cleanly 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.

---

CoreML stateful attention primitives

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.

Platform requirement: MLState and stateful CoreML models require iOS 18+ / macOS 15+ and Xcode 16+. Confirm your deployment target before adopting this pattern.

Illustrative implementation — adapt shape dimensions to your model’s sequence and vocabulary sizes:

// 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 calls. Each chunk adds to the cache without re-processing prior tokens — exactly the structure Flash Attention’s tiling requires.

---

Performance trade-offs: chunk size vs. latency

Chunk size is a tuning exercise, not a fixed answer. 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. That said, profile your specific model with Instruments’ Core ML template before committing — the table is a starting point, not a guarantee.

---

Quantization is a prerequisite, not an afterthought

No amount of attention tiling compensates for a model that doesn’t 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.

In my experience building production on-device inference pipelines, teams that retrofit quantization late consistently discover ANE fallback to CPU for layers the compiler can’t schedule. Profile with xcrun coremlcompiler compile and inspect the compute unit assignments before shipping.

---

Putting it together

Start at 256 tokens per chunk and 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. 4-bit weight quantization via coremltools is a prerequisite for fitting modern LLMs inside ANE working memory. Without it, tiling strategy is irrelevant — the model won’t schedule to ANE at all.

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.

---

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

---

Tags: ios mobile architecture swift