Principle:Ollama Ollama MLXRunner Architecture
| Knowledge Sources | |
|---|---|
| Domains | MLX, Apple Silicon |
| Last Updated | 2025-02-15 00:00 GMT |
Overview
The MLX Runner Architecture provides a native inference backend for running large language models on Apple Silicon hardware using Apple's MLX framework, enabling efficient GPU-accelerated inference through the Metal compute shaders available on M-series chips.
Core Concepts
MLX Framework Integration
MLX is an array computation framework designed specifically for Apple Silicon. Unlike GGML which targets a broad range of hardware, MLX leverages the unified memory architecture of Apple's M-series processors, allowing tensors to be shared between CPU and GPU without explicit data transfers. The MLX runner wraps this framework to provide LLM inference capabilities within Ollama's runner abstraction.
Runner Lifecycle
The MLX runner follows a client-server architecture. The server component initializes the MLX backend, loads model weights into unified memory, and exposes an inference endpoint. The client component communicates with the server to submit inference requests and stream token results back. This design mirrors the llama.cpp runner pattern, ensuring both backends can be managed uniformly by Ollama's scheduler.
Model Pipeline
The MLX runner implements a complete inference pipeline including model loading from GGUF or safetensors format, KV-cache management optimized for Metal, and token sampling. The pipeline component orchestrates the flow from input tokenization through forward pass execution to output token generation, leveraging MLX's lazy evaluation and just-in-time compilation for optimal throughput.
Cache Management
MLX's KV-cache implementation takes advantage of unified memory to avoid the CPU-GPU synchronization overhead present in discrete GPU systems. Cache entries are allocated in Metal-accessible memory and can be read or written by both the CPU and GPU execution units without explicit copy operations.
Implementation Notes
In the Ollama codebase, the MLX runner resides under x/mlxrunner/. The runner is structured with a server (server.go) that manages the MLX backend lifecycle, a runner (runner.go) that implements the common runner interface, a pipeline (pipeline.go) that orchestrates inference, and a client (client.go) that handles communication. The sampling subsystem under x/mlxrunner/sample/ provides MLX-native token sampling, and model definitions under x/mlxrunner/model/ provide architecture-specific forward pass implementations using MLX operations.