[Tech Review] Core AI: Redesigning the Conversion Stack for On-Device Generative Models

Hello, I'm Seongjae Park from OptAI.

This review comes from our team analysis of on-device LLM deployment pipelines. Running generative models such as LLMs and Diffusion models on a device differs from deploying conventional Vision/CNN models. Mutable KV cache state must be updated for every token, execution switches between prefill and decode, and outputs over a large vocabulary must also be handled.

Core AI is the on-device generative AI framework Apple announced at WWDC 2026, an official successor to Core ML for neural-network and transformer workloads. It opens the inference technology behind Apple Intelligence to third-party developers, bringing CPU/GPU/Neural Engine under one Swift API with AOT compilation and device-specific specialization. PyTorch models are converted to .aimodel through coreai-torch, while the Apple coreai-models repository provides export recipes for major models including Qwen and Mistral. Using Qwen3 as an example, this review compares its structure with the coremltools-based ANEMLL pipeline along three axes.

1. Introduction: The Two Pipelines

ANEMLL reimplements the HF model for ANE, Apple Neural Engine, then produces .mlmodelc through torch.jit.trace and coremltools.

HF checkpoint
   │  Load weights
   ▼
QwenForCausalLM  (anemll/models/qwen_model.py)   ← Reimplement the HF model for ANE
   │    - All Linear layers → nn.Conv2d(kernel=1), input [B, hidden, 1, seq]
   │    - KV cache = register_buffer("kv_cache_0", ...)  (Single unified K/V buffer)
   │    - forward inputs: input_ids, position_ids, causal_mask, current_pos, update_mask
   ▼
torch.jit.trace(wrapper, sample_inputs)          ← Capture a static graph
   ▼
ct.convert(traced,
           inputs=[...], outputs=[logits1..16],   ← Large vocab → 16 output partitions
           states=[ct.StateType(... kv_cache_0 ...)],
           compute_units=CPU_AND_NE,
           minimum_deployment_target=iOS18,
           convert_to="mlprogram")                ← Convert to coremltools MIL
   ▼
cto.coreml.palettize_weights(...)                ← LUT4/8 quantization (after conversion)
   ▼
.mlpackage ──compile──► .mlmodelc  (+ dedup, combine chunks)
   ▼
chat.py / Swift:  state = model.make_state();  model.predict(inputs, state)

Core AI extracts the graph with torch.export and produces .aimodel through its own MLIR-based compiler. Compression, including quantization and palettization, is integrated into the pipeline, with the tokenizer and metadata bundled into a single asset.

HF checkpoint
   │  AutoConfig → identify model_type → resolve the model class from the registry
   ▼
Qwen3 (coreai_models/models/ios/qwen3.py)        ← Own implementation using only HF weights
   │
   ▼ Compression integrated into the pipeline:
   │    macOS: QuantizerConfig + C4 calibration
   │    iOS  : KMeansPalettizerConfig
   ▼
torch.export.export(model.extend / .prompt_opt / .gather_embeddings / .load_embeddings,
                    dynamic_shapes={... torch.export.Dim ...})  ← 4 ExportedPrograms
   ▼
TorchConverter()
   .add_exported_program(... entrypoint_name=..., state_names=["key_cache","value_cache"])
   .to_coreai()
   .set_static_shape_config({(cache_len, q_len) combinations})           ← enumerated specialization
   .set_hardware_constraints(IOSurface, interleave=8, BC1S)     ← Declare the memory layout
   .optimize()                                                  ← MLIR optimization passes
   ▼
AIProgram → .aimodel
   ▼
bundle_llm_asset(): tokenizer/ + Automatically include metadata.json
   ▼
Load and generate with the Swift CoreAILanguageModels runtime

ANEMLL manually reshapes the graph for the hardware before passing it to the compiler. Core AI keeps model code conventional and delegates hardware specialization to the converter and compiler. This difference in philosophy underlies the design differences that follow.

2. Key Difference 1: Graph Capture — torch.jit.trace vs torch.export

The two stacks differ from the way they capture the PyTorch graph.

The TorchScript IR produced by trace is a C++ torch::jit::Graph object. It uses SSA Value/Node/Block structures, mixing aten and prim operations.

graph(%self : __torch__.Model,
      %x : Float(1, 10)):
  %w  : Tensor = prim::GetAttr[name="weight"](%self.linear)   # Read weights from module attributes
  %b  : Tensor = prim::GetAttr[name="bias"](%self.linear)
  %1  : Tensor = aten::linear(%x, %w, %b)                     # Keep the high-level aten op
  %d  : int    = prim::Constant[value=-1]()                   # Constants are prim nodes too
  %2  : Tensor = aten::softmax(%1, %d, ...)
  return (%2)

By contrast, the ExportedProgram produced by torch.export is based on an ATen-level FX graph. Functionalization removes in-place operations and mutation, normalizing it into a functional graph with weights promoted to graph inputs.

class GraphModule(torch.nn.Module):
    def forward(self, p_linear_weight, p_linear_bias, x):   # Weights are promoted to inputs
        permute  = torch.ops.aten.permute.default(p_linear_weight, [1, 0])
        addmm    = torch.ops.aten.addmm.default(p_linear_bias, x, permute)
        softmax  = torch.ops.aten._softmax.default(addmm, -1, False)
        return (softmax,)

# Graph signature:
#   inputs  = [p_linear_weight (PARAMETER), p_linear_bias (PARAMETER), x (USER_INPUT)]
#   outputs = [softmax (USER_OUTPUT)]

Control-flow handling also differs. Since trace records only the executed path, a data-dependent conditional leaves only the branch taken by the example input in the graph.

def forward(self, x):
    if x.sum() > 0:      # trace records only the branch taken by the example input
        return x * 2
    else:
        return x * 3

Export detects such branches and raises an error or requires an explicit control-flow operator such as torch.cond. Rather than silently generating an incorrect graph, it exposes the issue during conversion.

An ExportedProgram is not merely a graph: it is a container combining graph_module, graph_signature, state_dict, and range_constraints. It contains decomposed, functional ATen operations with weights promoted to inputs, providing regular, predictable input for a converter to traverse node by node and map to IR.

Using this regular structure, TorchConverter lowers ATen operations to Core AI operations. Lowering rules are defined in coreai_torch/_aten_to_core.py.

def replace_softmax(
    values_map: dict[str, Value], node: fx.Node, loc: Location
) -> Value:
    x = _get_operand(values_map, node, 0)
    dim = node.args[1]
    return coreai.softmax(x, dim + x.type.rank if dim < 0 else dim)

The conversion flow can be summarized as follows.

to_coreai()
  ↓
_get_graph_op()
  ↓
for node in exported_program.graph_module.graph.nodes
  ↓
_get_operation(node)
  ↓
_handle_call_function_op(node)
  ↓
namespace == "aten"
  ↓
_aten_to_core_resolver[target](values_map, node, loc)
  ↓
Create coreai.*Op
  ↓
_values_map[node.name] = op_result

3. Key Difference 2: KV Cache — Model Ownership vs Converter Promotion

In ANEMLL, the model owns the cache. Cache tensors become part of the model through register_buffer, with K and V combined in one buffer shaped (num_layers*2, n_kv_heads, state_length, head_dim).

# anemll/models/qwen_model.py — Register a single unified K/V buffer
cache_size = (
    2 * config.num_hidden_layers,
    config.num_key_value_heads,
    config.state_length,
    self.head_dim
)
self.register_buffer("kv_cache_0", torch.zeros(cache_size, dtype=MODEL_DTYPE, device=TEST_DEVICE))

Tracing captures in-place slice assignments, which become ct.StateType. At inference time, the runtime manages the cache implicitly through make_state() → predict(inputs, state). Separate state_length, the physical length, and context_length, the logical length, handle sliding and rotation.

In Core AI, the converter promotes inputs to state. Model code treats the cache as ordinary input/output tensors, key_cache → new_k_cache, while state declarations and memory layout are specified declaratively during conversion.

# coreai_models/export/ios.py
state_names = ["key_cache", "value_cache"]
converter.add_exported_program(extend_program, input_names=..., state_names=state_names, ...)

cache_constraints = HardwareConstraints(
    AllocationType.IOSurface,
    interleave=[1, 1, KV_CACHE_INTERLEAVE_FACTOR, 1, 1],   # ANE-friendly alignment
    alignments=[1, 1, 1, 1, KV_CACHE_INTERLEAVE_FACTOR * max_context_length, 1],
)

The cache shape is (num_layers, 1, kv_cached_embed_size, 1, cache_len): a BC1S layout with sequence on the last axis.

Separate iOS and macOS paths. Core AI maintains different paths for these platforms. ANE was designed for CNN/Vision workloads and favors static shapes and a specific memory layout, B, C, 1, S, with a limited operation set and substantial dispatch overhead. The principal macOS target is GPU, Metal, whose general-purpose compute supports dynamic shapes and flexible layouts. Code that is simultaneously ANE-friendly and GPU-friendly is effectively impractical for the same LLM. iOS therefore uses fixed cache length plus enumerated specialization, while macOS uses dynamic sequence length through torch.export.Dim("k_seq_len", ...).

This introduces a problem. torch.export functionalizes the graph to remove in-place operations, but KV cache updates inherently overwrite a specific slice of a buffer in place. Decomposing this into ordinary ATen slice and copy operations prevents the compiler from recognizing the intent of a state update.

Core AI solves this with a custom op.

@torch.library.custom_op("coreai::mutable_slice_update", mutates_args=["x"])
def mutable_slice_update(
    x: Tensor,
    update: Tensor,
    begin: Tensor,
    end: Tensor,
) -> Tensor:
    """
    Mutable slice update operation for cache updates.
    Updates a slice of tensor x with the update tensor using dynamic begin/end indices.
    """
    begin = torch.split(begin, 1, dim=0)
    end = torch.split(end, 1, dim=0)
    slices = tuple(slice(b.item(), e.item()) for b, e in zip(begin, end, strict=False))
    x[slices] = update
    # Note: Not actually in-place for torch
    return x.clone()

In registration, mutates_args=["x"] declares that the operation has the side effect of modifying input x. register_fake provides a fake, or meta, implementation that reports output shape/dtype without actual computation.

Register the op with the PyTorch dispatcher through @torch.library.custom_op
  ↓
Register the eager-execution implementation
  ↓
Register the fake/meta implementation through @register_fake
  ↓
Enable shape inference during torch.export / tracing
  ↓
The Core AI converter recognizes this op as a state update primitive / custom lowering

The key is to preserve the KV cache update as the single graph node coreai::mutable_slice_update(x, update, begin, end). The compiler interprets it as a slice update to a mutable state buffer rather than a simple slice/copy, lowering it to a state-update primitive or a target-specific kernel. Core AI documentation also describes model authoring through composite operations, custom-op lowering, and inline Metal GPU kernels, alongside conversion of existing PyTorch models.

Placing the Qwen3 conversion code side by side makes the difference in how the two stacks express the same KV cache update clear.

As shown earlier, ANEMLL performs in-place slice assignments into a unified buffer in the model. During conversion, GetTransformerStates in qwen_converter.py declares that buffer as a StateType so coremltools recognizes it as state.

# anemll/models/qwen_model.py — In-place slice updates during prefill
key_idx = layer_in_group_idx
value_idx = layer_in_group_idx + layers_per_group

# Store the full sequence length in prefill mode
seq_length = key_states.shape[2]  # Get actual sequence length
kv_cache[key_idx:key_idx + 1, :, current_pos:current_pos + seq_length, :] = key_states
kv_cache[value_idx:value_idx + 1, :, current_pos:current_pos + seq_length, :] = value_states

# anemll/ane_converter/qwen_converter.py — GetTransformerStates
num_layers_this_part = num_layers * 2   # 2 times the layer count because K+V are unified

states = [
    ct.StateType(
        wrapped_type=ct.TensorType(
            shape=(
                num_layers_this_part,
                model.config.num_key_value_heads,
                model.config.state_length,
                head_dim,
            ),
            dtype=np.float16,
        ),
        name=f"{prefix}kv_cache_0",  # Only one group for unified cache
    )
]

Core AI, coreai_models/models/ios/qwen3.py, expresses mutation through a custom op. After torch.export functionalization, the conversion stage declares it as state. Dynamic begin/end tensors specify the slice position, allowing arbitrary-position updates without rebuilding the graph.

begin, end = self.gen_slice_args(layer_idx, offset, num_token_updates)
# update k — iOS updates dimension 4 (the last dimension)
mutable_slice_update(x=self._k_cache, update=k.unsqueeze(0), begin=begin, end=end)
# update v
mutable_slice_update(x=self._v_cache, update=v.unsqueeze(0), begin=begin, end=end)

4. Key Difference 3: Compiler Stack — MIL vs MLIR

[ANEMLL]
TorchScript → coremltools MIL (its own IR, ≠MLIR) → MIL optimization passes → mlprogram(.mlmodelc)
                                                                → Apple private ANE compiler (espresso)

[Core AI]
ExportedProgram → MLIR(coreai dialect) → MLIR optimization passes(optimize) → AIProgram(.aimodel)
                                                                → Apple private ANE runtime/compiler

Core AI uses MLIR as its IR. export/mlir_ops.py imports its own MLIR dialect through from coreai._compiler.dialects import coreai, coreaix, and performs .to_mlir() lowering as well as post-MLIR INT4 quantization.

The philosophical difference is clearest here. ANEMLL manually reshapes the PyTorch graph for ANE before compilation: 1×1 Conv2d, per-head SDPA, logits split into 16 outputs, and a register_buffer cache. Reducing compiler freedom gives more control and predictability, validated with anemll-profile. Core AI preserves patterns such as RoPE/RMSNorm/SDPA as named composites when passing them to MLIR, delegating fusion and layout decisions through MLIR optimization passes and HardwareConstraints. This moves toward greater automation and a higher optimization ceiling.

5. Graph Partitioning and Prefill/Decode

LLM inference consists of prefill, processing the full prompt at once, and decode, generating tokens one at a time. The stacks represent these stages differently.

The 4 entrypoints produced by Core AI iOS export handle load_embeddings, returning the embedding table; gather_embeddings, token ID → embedding; extend, the decode-mode forward pass; and prompt_opt, the prefill-mode forward pass, respectively, in export/ios.py. ANEMLL splits the model across files and removes duplicate weights with postprocessing scripts. Core AI exports four entrypoints from one module and bundles them into a single multifunction asset. Weight sharing is structural, and the deployment unit is simpler.

Large-vocabulary handling also differs. ANEMLL splits logits into 16 outputs to work around ANE constraints. Core AI keeps lm_head as a single nn.Linear(hidden_size, vocab_size) layer but separates the large vocab_size × hidden_size embedding table from the transformer graph. load_embeddings returns the int8-stored table, and gather_embeddings uses the coreai::fused_dequant_gather_reshape custom op to fuse dequantization, gather, and reshape, extracting only the required token embeddings. With tie_word_embeddings enabled, this table is also shared with lm_head.

# coreai_models/primitives/ios/embedding.py — Fused custom op for separating embeddings
@torch.library.custom_op("coreai::fused_dequant_gather_reshape", mutates_args=[])
def fused_dequant_gather_reshape(
    embedding_table: torch.Tensor,
    input_ids: torch.Tensor,
    scale: torch.Tensor,
    final_shape: list[int],
) -> torch.Tensor:
    return (embedding_table[input_ids].to(scale.dtype) * scale).reshape(final_shape)

6. Conclusion: From Manual Preparation to Delegation

The difference can be summarized in one sentence: ANEMLL manually adapts the graph to the hardware, while Core AI keeps the model conventional and delegates adaptation to the converter and compiler. This philosophy consistently shapes graph capture, trace versus export; KV cache management, model ownership versus converter promotion; the compiler, manual MIL preparation versus MLIR delegation; and deployment units, multiple files versus a single multifunction asset.

Regular torch.export graphs, preservation of state intent through custom ops, MLIR dialects, and declarative hardware constraints: Core AI redesigns the full conversion pipeline for on-device generative models. The next article will cover benchmark measurements from converting and running an actual model on a device.

The container produced by torch.export, combining a functionalized ATen-level FX graph with graph_signature, state_dict, and range_constraints. It provides regular, predictable input to a converter.

Declaring a KV cache, treated as ordinary input/output tensors in model code, as runtime-managed mutable state during conversion. This is central to separating model code from hardware specialization.

A memory layout of Batch, Channel, 1, Sequence, with the sequence on the last axis, favored by ANE.

For hardware without dynamic-shape support, enumerate the permitted cache_len, q_len combinations to include specialized graphs within a single model.

References

1. Apple, "Integrate on-device AI models into your app using Core AI" — WWDC26 Session 326 (developer.apple.com/videos/play/wwdc2026/326/)
2. InfoQ, "Apple Launches Core AI for Apple-Silicon Optimized On-Device Generative AI" (2026.06)
3. apple/coreai-models — official Core AI models and export recipes (github.com/apple/coreai-models)
4. ANEMLL — open-source LLM conversion project for Apple Neural Engine (github.com/anemll/anemll)
5. PyTorch, torch.export documentation — ExportedProgram, Dim, torch.cond
6. PyTorch, torch.library documentation — custom_op, register_fake
7. coremltools documentation — MIL, StateType, palettize_weights