iOS: supplying your own [KVCache] to generate() terminates the process with a larger model (macOS unaffected; MLX-managed cache unaffected)

[srbryers]

What happens

On iOS, passing your own [KVCache] to
generate(input:cache:parameters:context:) terminates the process during
prompt prefill when using a larger model. The same code path completes on
macOS, and completes on iOS when MLX creates the cache itself.

There is no crash report, no NSException, and no jetsam event.
os_proc_available_memory() reports ~1.1 GB still available on the log
line immediately before the process disappears.

Measured

iPhone 16 Pro (8 GB), iOS 26, mlx-swift-examples 2.29.1, mlx-swift 0.29.1.
Fresh process per cell — relaunch, run one scenario, nothing before it.
4,009-token prefill, 8 tokens generated, kvBits = 8 in every case.

cache	Qwen3-1.7B-4bit	Qwen3-4B-Instruct-2507-4bit-DWQ-2510
MLX-managed (cache: nil)	completes	completes — 347 MB, 24.7 s
supplied KVCacheSimple	completes — 713 MB	terminates at 1,098 MB free
supplied QuantizedKVCache	completes — 311 MB	terminates at 1,111 MB free

macOS (M-series, same package versions) completes every cell including the
4B.

Memory does not seem to explain it. The supplied quantized cache measures
311 MB on the 1.7B, so roughly 490 MB on the 4B, against 1,111 MB of
reported headroom — and the managed cache on the same model and prompt
completes using 347 MB.

MLX.GPU.set(cacheLimit: 32 * 1024 * 1024) is called before loading in all
cases. Without it, even the managed scenario terminates, which is a
separate and much more obvious failure.

Reproduction

MLX.GPU.set(cacheLimit: 32 * 1024 * 1024)

let container = try await LLMModelFactory.shared.loadContainer(
    configuration: ModelConfiguration(id: "mlx-community/Qwen3-4B-Instruct-2507-4bit-DWQ-2510"))

try await container.perform { (context: ModelContext) in
    var parameters = GenerateParameters(temperature: 0.0)
    parameters.maxTokens = 8
    parameters.kvBits = 8
    parameters.kvGroupSize = 64
    parameters.quantizedKVStart = 0

    // ~4,000 tokens of ordinary text
    let filler = String(repeating: "The quick brown fox jumps over the lazy dog. ", count: 400)
    let input = try await context.processor.prepare(input: .init(chat: [.user(filler)]))

    // nil            -> completes
    // supplied cache -> process terminates during prefill, on iOS, 4B only
    let cache: [KVCache]? = context.model.newCache(parameters: nil)

    for await _ in try MLXLMCommon.generate(
        input: input, cache: cache, parameters: parameters, context: context
    ) {}
}

Swapping cache between nil and context.model.newCache(parameters: nil)
is the only change between the passing and failing runs.

Why supply a cache at all

An app whose prompt is largely constant between turns — system prompt plus
tool schemas — can keep the KV state and re-prefill only what changed. Here
that is 3,911 of 4,151 tokens identical turn to turn, and re-prefilling
them costs 25–28 seconds per message on device. Reuse measured 2.5x faster
on macOS (522 s → 211 s over a 140-case suite) with an identical score, so
it is worth quite a lot if it can be made to work on iOS.

What I have not established

Any root cause. I noticed QuantizedKVCache.update(keys:values:) is a
fatalError (KVCache.swift:743) and initially assumed that was it, but the
supplied-quantized scenario completes on the 1.7B and on macOS, so that
cannot be the whole story — and suppliedSimple fails identically without
involving that type at all.

Happy to run further instrumentation on the device if it would help.

[davidkoski]

Are you using the "Increased Memory Limit" entitlement?

It is curious that you passing the KVCache vs not is making a difference -- under the hood the code would just allocate one and discard at the end.