Use prefill for long prompts

Two jobs, two functions

A language model does two different jobs. Prefill reads the whole prompt and fills the KV cache. Decode makes one token (or a few) at a time. The first token waits for prefill, so for long prompts prefill is most of the wait.

Apple's export recipes in apple/coreai-models reflect this:

• Many macOS language model recipes (for example qwen3.py, gemma3_text.py, mistral.py) set exports_prefill_graph = True, with the comment "Emit a second, prefill-only prefill entrypoint beside main."
• The Neural Engine rules in the repo list static entrypoints per shape: extend_{ctx}_{len} for token generation (logits plus updated KV) and prompt_opt_{ctx}_{len} for fast prefill (KV only, no logits).
• The Gemma 4 iOS export ships prefill (prompt_opt) at query length 64 and decode (extend) at query length 8. Its comment says the runner prefills in 64-token chunks while more than 64 tokens remain, and runs the rest as 8-token extend steps. It also uses a small set of context buckets (1024, 8192, 32768, 131072) to stay under the accelerator's per-program I/O cap.

See what your export contains

Function names depend on the recipe and platform. Do not guess them. Open the .aimodel in Xcode and look at the Functions tab, or list them in code with functionNames:

import CoreAI

func describeFunctions(_ model: AIModel) {
    for name in model.functionNames.sorted() {
        print(name)
    }
}

Load both functions before the first request

Apple's docs say loading a function prepares the resources to run it and "can also be expensive". If you load the prefill function only when a long prompt arrives, the user pays that cost on the first token. Load the functions you will need while the app is idle.

import CoreAI

struct LMFunctions {
    let prefill: InferenceFunction?
    let decode: InferenceFunction
}

func loadLMFunctions(_ model: AIModel, prefillName: String?, decodeName: String) throws -> LMFunctions {
    guard let decode = try model.loadFunction(named: decodeName) else {
        throw CocoaError(.featureUnsupported)
    }
    let prefill = try prefillName.flatMap { try model.loadFunction(named: $0) }
    return LMFunctions(prefill: prefill, decode: decode)
}

/// Same split the coreai-models Gemma 4 iOS export describes: 64-token prefill chunks, small decode steps for the rest.
func plan(promptTokens: Int, prefillChunk: Int = 64) -> (prefillChunks: Int, tailTokens: Int) {
    let chunks = promptTokens / prefillChunk
    return (chunks, promptTokens - chunks * prefillChunk)
}

The actual inputs and state tensors (token ids, positions, KV cache) differ per recipe. Read them from function.descriptor or the Functions tab. If you run language models through LanguageModelSession instead of raw functions, Apple's Running a Core AI model in a Foundation Models session guide and the coreai-models Swift runtime handle this selection for you. The repo's tests show the runtime choosing prompt_opt_ graphs for prefill.

Dynamic shapes: expectFrequentReshapes

By default, Core AI optimizes a dynamic-shape function for each new input shape it sees. Apple says that for a dynamic-shape LLM whose sequence grows one token at a time, this can cost more than it saves. Set expectFrequentReshapes to true so Core AI runs the generic dynamic version:

import CoreAI

func dynamicLMOptions() -> SpecializationOptions {
    var options = SpecializationOptions.default
    options.expectFrequentReshapes = true
    return options
}

Remember from lesson 02: these options are part of the cache key. Use the same value when you call specialize.

Sources

• Apple: SpecializationOptions.expectFrequentReshapes
• Apple: AIModel.functionNames
• apple/coreai-models repository