When a machine learning graph needs to execute across multiple hardware backends—CPU, GPU, Neural Engine, or specialized accelerators—the runtime must decide which backend executes which operations. A naïve approach assigns all operations to a single backend, but real hardware has uneven capabilities: the Apple Neural Engine excels at dense matrix multiplications but doesn’t support all activation functions; MLX handles GPU-based tensor ops but has different precision constraints; BLAS provides reliable CPU fallbacks for everything but at lower throughput.
Capability-based backend partitioning solves this by having each backend declare exactly which operations it supports, at what precision levels, and for which tensor shapes. The partitioning algorithm then walks the computation graph and assigns each node to the most capable backend, inserting data transfer operations at partition boundaries where tensors must move between backends (e.g., from Neural Engine memory to CPU memory).
Complementing partitioning is the graph cache: an LRU cache that stores compiled executors, weight tensors, and analysis metadata keyed by graph structure and input shapes. In autoregressive decoding, where the same graph structure executes thousands of times with only the sequence position changing, the graph cache eliminates redundant compilation and weight loading, reducing per-token latency from milliseconds of compilation overhead to microseconds of cache lookup.
Each backend registers a capability set: the operations it supports, the data types it handles (f16, f32, bf16), and any shape constraints (e.g., “matmul only for matrices with dimensions divisible by 16”). The partitioner queries these capabilities when deciding where to place each graph node. If a backend doesn’t declare support for an operation, that node cannot be assigned to it.
Partitioning splits the computation graph into subgraphs, each assigned to a single backend. The algorithm balances two competing goals: maximize use of fast backends (Neural Engine, GPU) and minimize cross-backend data transfers, which are expensive because they involve copying tensors between different memory spaces. A greedy algorithm assigns each node to the fastest capable backend, then merges adjacent nodes on the same backend into contiguous subgraphs.
When an edge in the graph crosses a partition boundary (node A runs on CoreML, node B runs on BLAS), the runtime inserts a transfer node that copies the output tensor from A’s memory space to B’s memory space. On Apple Silicon, this may involve IOSurface-backed shared memory to avoid actual copies, but the transfer node still handles format conversion (e.g., CoreML’s NHWC layout to BLAS’s row-major layout).
The graph cache maps a cache key (graph topology hash + input shapes + backend assignment) to a compiled executor. On cache hit, the runtime skips tracing, optimization, and partitioning entirely, jumping straight to execution with pre-compiled kernels and pre-loaded weights. The cache uses LRU eviction to bound memory usage, evicting the least recently used compiled graphs when the cache reaches its size limit.
In autoregressive text generation, each token decode uses the same graph structure but with an incremented sequence position. The graph cache recognizes this pattern and reuses the compiled executor across decode steps, only updating the position-dependent inputs (KV cache pointers, attention mask). This turns what would be thousands of compilations into a single compilation plus thousands of cache hits.
Each backend registers its capabilities at initialization time. The capability set is a compile-time constant for each backend, making the partitioner’s queries zero-cost.
const BackendCapability = struct {
supported_ops: EnumSet(Op),
supported_dtypes: EnumSet(DType),
max_tensor_dims: u8,
requires_aligned_shapes: bool,
alignment: u32, // e.g., 16 for CoreML
priority: u8, // higher = preferred
};
const coreml_caps = BackendCapability{
.supported_ops = EnumSet(Op).initMany(&.{
.matmul, .add, .relu, .softmax, .layer_norm,
.conv2d, .gelu, .silu,
}),
.supported_dtypes = EnumSet(DType).initMany(&.{ .f16, .f32 }),
.max_tensor_dims = 5,
.requires_aligned_shapes = true,
.alignment = 16,
.priority = 3, // highest priority
};
const blas_caps = BackendCapability{
.supported_ops = EnumSet(Op).initAll(), // supports everything
.supported_dtypes = EnumSet(DType).initAll(),
.max_tensor_dims = 8,
.requires_aligned_shapes = false,
.alignment = 1,
.priority = 1, // fallback
};The partitioner walks the graph in topological order, assigning each node to the highest-priority backend that supports it. Adjacent nodes on the same backend are merged into subgraphs. At partition boundaries, transfer nodes are inserted.
fn partitionGraph(
graph: *Graph,
backends: []const BackendCapability,
) []Partition {
var assignments = allocator.alloc(BackendId, graph.nodes.len);
// Assign each node to the best capable backend
for (graph.nodes.items(), 0..) |node, i| {
assignments[i] = selectBackend(node, backends);
}
// Merge adjacent same-backend nodes into partitions
var partitions = ArrayList(Partition).init(allocator);
var current = Partition{
.backend = assignments[0],
.nodes = ArrayList(*Node).init(allocator),
};
for (graph.topologicalOrder()) |node_idx| {
if (assignments[node_idx] != current.backend) {
partitions.append(current);
current = Partition{
.backend = assignments[node_idx],
.nodes = ArrayList(*Node).init(allocator),
};
}
current.nodes.append(&graph.nodes.items()[node_idx]);
}
partitions.append(current);
// Insert transfer nodes at partition boundaries
insertTransferNodes(&partitions);
return partitions.toOwnedSlice();
}
fn selectBackend(node: *Node, backends: []const BackendCapability) BackendId {
var best: ?BackendId = null;
var best_priority: u8 = 0;
for (backends, 0..) |cap, id| {
if (cap.supported_ops.contains(node.op) and
cap.supported_dtypes.contains(node.dtype) and
(!cap.requires_aligned_shapes or isAligned(node.shape, cap.alignment)))
{
if (cap.priority > best_priority) {
best = @intCast(id);
best_priority = cap.priority;
}
}
}
return best orelse @panic("no backend supports this operation");
}The graph cache uses a hash of the graph structure (operation types, edge connectivity, tensor shapes) as the cache key. On cache hit, it returns a pre-compiled executor. On miss, it compiles the partitioned graph and stores the result.
const GraphCache = struct {
entries: LRUMap(u64, CachedExecutor),
max_entries: usize,
weight_store: WeightStore,
fn getOrCompile(
self: *GraphCache,
graph: *Graph,
partitions: []Partition,
) *Executor {
const key = self.computeCacheKey(graph);
if (self.entries.get(key)) |cached| {
cached.last_used = std.time.timestamp();
return cached.executor;
}
// Cache miss: compile the partitioned graph
const executor = compilePartitions(partitions, self.weight_store);
if (self.entries.count() >= self.max_entries) {
self.evictLRU();
}
self.entries.put(key, CachedExecutor{
.executor = executor,
.last_used = std.time.timestamp(),
.graph_hash = key,
});
return executor;
}
fn computeCacheKey(self: *GraphCache, graph: *Graph) u64 {
var hasher = std.hash.Wyhash.init(0);
for (graph.nodes.items()) |node| {
hasher.update(std.mem.asBytes(&node.op));
hasher.update(std.mem.asBytes(&node.shape));
hasher.update(std.mem.asBytes(&node.dtype));
for (node.inputs[0..node.num_inputs]) |input| {
hasher.update(std.mem.asBytes(&input.index));
}
}
return hasher.final();
}
};During text generation, the decode loop calls the graph cache on every token. The first call compiles the graph; subsequent calls get cache hits because the graph topology is identical—only the input values (token ID, KV cache state) change between iterations.
fn generateTokens(model: *Model, prompt: []const u32, max_tokens: usize) []u32 {
var tokens = ArrayList(u32).init(allocator);
tokens.appendSlice(prompt);
var kv_cache = KVCache.init(model.config);
for (0..max_tokens) |_| {
// Trace the graph (same structure every iteration)
const graph = model.traceForward(tokens.last(), &kv_cache);
// Optimize (cached after first call)
runPassPipeline(graph);
// Partition and get compiled executor (cached after first call)
const partitions = partitionGraph(graph, backends);
const executor = graph_cache.getOrCompile(graph, partitions);
// Execute - the only work done on cache hits
const logits = executor.run();
const next_token = sampleToken(logits);
if (next_token == EOS_TOKEN) break;
tokens.append(next_token);
}
return tokens.toOwnedSlice();
}Large model weights (billions of parameters) shouldn’t be duplicated across cached executors. The weight store holds a single copy of each weight tensor, and cached executors reference weights by pointer. When a cached executor is evicted, only the compiled kernels are freed—the weights remain as long as any executor references them.
The initial capability-based assignment is a heuristic. For production workloads, the runtime can profile actual execution times per backend and adjust partitioning. If CoreML’s matmul is slower than expected for a specific shape (due to alignment overhead), the profiler can override the capability-based assignment and route that node to MLX instead.
When a subgraph on backend A finishes and its output needs to transfer to backend B, the transfer can happen concurrently with backend A starting on its next subgraph (if there is one). This pipelining hides transfer latency behind compute, especially effective when the GPU and Neural Engine operate on non-overlapping graph segments.