WWDC Quick Look đź’“ By SwiftGGTeam
中文 EN 日本語
Deploy machine learning and AI models on-device with Core ML

Deploy machine learning and AI models on-device with Core ML

WWDC 2024
Watch original video

Highlight

Core ML adds the MLTensor type for efficient tensor computation, the State API optimizes inference efficiency for stateful models, and multi-function models let a single MLPackage host multiple capabilities.

Core Content

When running ML models on device, developers face a practical reality: model inference is only part of the workflow. For large language models, generating text requires the decoder to sample the next token from vocabulary scores—this “glue code” including probability distribution computation, temperature adjustment, and Top-K sampling previously required hand-written implementations or low-level APIs, with lots of code and error-prone logic.

Core ML’s new MLTensor addresses this. It’s a multi-dimensional array type offering NumPy/PyTorch-like math operations and transforms, accelerated on Apple silicon’s CPU, GPU, and Neural Engine. The session demo shows the same decoding logic implemented with MLTensor uses far less code than hand-written low-level API code.

Another pain point is stateful inference with KV Cache management. Previously, manual cache management required passing the cache as input each iteration and reading updated values from output—data shuttled back and forth between memory and compute units. Core ML’s new State API handles state management internally, reducing data movement overhead. In the session, Mistral 7B on an M3 Max MacBook Pro completed in about 5 seconds with State versus about 8 seconds without—roughly 1.6x faster (12:11).

Deploying multiple related models also creates redundancy—for example, one diffusion model with different style adapters previously required deploying multiple independent models. Multi-function models let a single MLPackage contain multiple function entry points, sharing base weights and reducing deployment size and memory usage.

Additionally, iOS 18’s inference stack received low-level optimizations—many models get performance gains without recompilation or code changes (02:27).

Detailed Content

MLTensor: on-device tensor computation

MLTensor is Core ML’s new multi-dimensional array type, defined by shape and scalar type, creatable from MLShapedArray or nested scalar collections (04:03):

// Create from MLShapedArray
let tensor1 = MLTensor(shapedArray: shapedArray)

// Create from a nested collection of scalars
let tensor2 = MLTensor([[1.0, 2.0], [3.0, 4.0]])

Key points:

  • MLTensor can be constructed from MLShapedArray or nested literals
  • Like MLShapedArray, described by shape and scalar type
  • Nesting depth of literals determines dimension count

Tensors support element-wise math, automatic broadcasting, comparison, and masking (04:33):

// Element-wise addition and multiplication, with literals broadcast automatically
let result = (tensor1 + 10) * tensor2

// Calculate the mean
let mean = result.mean()

// Create a Boolean mask with a comparison operation
let mask = result .> mean

// Multiply by the mask to filter values
let filtered = mask * result

Key points:

  • Literals automatically broadcast to compatible shapes when operating with tensors
  • .mean() computes mean of all elements
  • .> comparison returns boolean tensor mask
  • Multiplying boolean mask with tensor zeros false positions

Indexing and shape transforms (04:57):

// Slice: get the first row of the matrix
let firstRow = matrix[0]

// Reshape
let reshaped = firstRow.reshape([1, -1])

Key points:

  • Index syntax slices by dimension, similar to Python numeric libraries
  • .reshape() changes shape without changing data

All tensor operations execute asynchronously—explicit materialize is required to access data (05:12):

let shapedArray = try await result.materialize()

Key points:

  • Async execution ensures upstream operations complete before data is available
  • materialize converts MLTensor back to MLShapedArray

Simplifying LLM decoding with MLTensor

The session compares decoders using HuggingFace Swift Transformer + Mistral 7B (06:33). Language models output scores (logits) for all vocabulary words; the decoder picks the next token—strategies include greedy decoding (highest score) and Top-K sampling (random sample among top K words), with temperature adjusting probability distribution flatness (06:57).

Comparison result: MLTensor implementation uses far less code for the same functionality than hand-written low-level API version. MLTensor suits common ML operations, but low-level APIs still have value when fine-grained control is needed (07:57).

State API: optimizing stateful inference

Traditional manual KV Cache management (11:04):

// Manually preallocate the cache
var keyValueCache = createEmptyCache()

for token in inputTokens {
    // Pass the cache as input
    let prediction = try model.prediction(from: input, using: keyValueCache)
    // Retrieve the updated cache from the output
    keyValueCache = prediction.cache
}

With State API (11:26):

// Create the model instance and preallocate state
let state = model.makeState()

for token in inputTokens {
    // Pass in the state; updates happen in place
    try model.prediction(from: input, state: state)
}

Key points:

  • model.makeState() pre-allocates state buffers by Core ML, returning a state handle
  • Access buffers and control state lifecycle through the handle
  • State updates in place—no reading from output and reassigning
  • Model must explicitly enable State support during preparation (coremltools conversion)

You can confirm State availability in Xcode model preview’s Predictions tab (10:48).

Multi-function model

The session uses Stable Diffusion XL + two style adapters (sticker and storybook) as an example (14:25). Both adapters share the same UNet base model; merged into one MLPackage, each exposes as one function:

// Specify the function name when loading
let model = try MLModel(
    contentsOf: url,
    configuration: MLModelConfiguration(functionName: "sticker")
)

Key points:

  • MLModelConfiguration(functionName:) specifies which function to load
  • Default function used when not specified
  • Different functions can have different input/output signatures
  • Xcode preview shows all available functions

The same pipeline can be reused for both functions since sticker and storybook have identical input/output signatures. Multiple models in a pipeline stitch seamlessly via MLTensor (14:35).

Performance tool enhancements

Core ML performance reports add two pieces of information (16:05):

  • Per-operation time estimate: estimated time = median prediction time Ă— operation’s estimated relative cost; sortable to find bottlenecks
  • Unsupported reason hints: hover over unsupported operations for reasons (such as unsupported data type), helping trace back to preparation phase changes

Reports now support export and cross-version comparison—evaluate model changes without writing code (16:56).

New MLComputePlan API provides programmatic access: model structure, each operation’s supported/preferred compute device, state support, estimated relative cost (17:13).

Core Takeaways

  • Replace hand-written decoding logic with MLTensor: On-device LLM token decoding (probability computation, sampling, temperature adjustment) is the most typical use case. Rewrite probability computation and filtering with MLTensor first, verify functional consistency, then gradually replace the entire decoder. Less code, and async execution avoids unnecessary synchronization waits.

  • Enable State API to reduce KV Cache overhead: If your model involves multi-iteration inference, re-export with coremltools and enable State. Mistral 7B measured 1.6x speedup—the larger the model and cache, the greater the benefit. Confirm State availability in Xcode preview first, then replace manual cache management with makeState().

  • Merge shared-weight models with multi-function models: When deploying multiple style adapters or task variants, merge them into a single MLPackage with each variant as one function. Reduces install size and runtime memory; specify functionName at load time to switch capabilities.

Comments

GitHub Issues · utterances