MLX

MLX builds computation graphs without immediate execution by deferring operations until explicit eval() calls. Operations create graph nodes in the array class that represent pending computations; the framework delays actual kernel dispatch to the backend until evaluation is triggered, enabling graph optimization and memory efficiency. This lazy model is implemented through a deferred execution pattern where each operation returns an array wrapping a computation node rather than executing immediately.

MLX abstracts hardware differences through a multi-backend system where the core API is platform-agnostic and each backend (Metal for Apple Silicon, CUDA for NVIDIA, CPU fallback) implements the same Primitive interface with eval_cpu(), eval_gpu(), and device-specific methods. The framework routes operations to the appropriate backend at runtime based on device selection, allowing identical Python code to run on M1/M2/M3/M4 chips, NVIDIA GPUs, or CPU without modification.

MLX-LM is a companion library for running large language models (LLMs) on Apple Silicon, providing model loading, tokenization, and text generation with support for popular architectures (Llama, Mistral, Phi, etc.). The library handles model quantization, prompt caching for efficient multi-turn conversations, and generation algorithms (greedy, beam search, top-k sampling). Models are loaded from Hugging Face Hub and automatically optimized for Apple Silicon.

MLX-VLM extends MLX-LM to support vision-language models (VLMs) that process both images and text, enabling tasks like image captioning, visual question answering, and image understanding. The library handles image preprocessing, vision encoder inference, and integration with language model components. Models like LLaVA and others are supported with automatic optimization for Apple Silicon.

MLX enables users to define custom primitives and kernels that integrate with the framework's operation system and autodiff. Custom primitives inherit from the Primitive class and implement eval_cpu() and eval_gpu() methods for different backends. Users can write Metal Shading Language (MSL) kernels for GPU computation or C++ code for CPU, and the framework automatically handles autodiff by requiring VJP/JVP definitions for custom operations.

MLX uses Nanobind to create efficient Python-C++ bindings that expose the C++ core to Python with minimal overhead. The bindings support NumPy-style indexing (slicing, fancy indexing, boolean indexing) on arrays, enabling Pythonic array manipulation. Nanobind generates type-safe bindings that preserve performance while providing a natural Python API.

MLX uses Nanobind (mlx/python/src) to create efficient Python-C++ bindings with minimal overhead. Nanobind generates type-safe bindings that preserve C++ semantics while exposing a Pythonic API. The binding layer handles array conversion, type promotion, and error propagation. Integration with lazy evaluation means Python operations return unevaluated computation graphs, enabling efficient batching and optimization.

MLX implements automatic differentiation through Vector-Jacobian Products (VJP) for reverse-mode autodiff and Jacobian-Vector Products (JVP) for forward-mode autodiff, building gradient computation graphs that mirror the forward computation. The framework traces operations to construct a computation graph, then applies the chain rule in reverse (for backprop) or forward (for forward-mode) to compute gradients. Both modes are composable and can be nested for higher-order derivatives.

MLX provides vmap (vectorization map) as a transform that automatically vectorizes scalar functions over batch dimensions without manual broadcasting or loop unrolling. vmap takes a function and a batch axis specification, then generates vectorized code that applies the function to each batch element in parallel. This is implemented as a transform in mlx/transforms.cpp that modifies the computation graph to add batch dimensions and broadcast operations accordingly.

MLX compiles computation graphs into optimized kernel sequences through a compilation system (mlx/compile.cpp) that fuses operations, eliminates redundant computations, and generates backend-specific code. The compiler analyzes the computation graph, identifies fusion opportunities (combining multiple operations into single kernels), and generates optimized code for Metal or CUDA. This happens transparently when eval() is called, reducing memory bandwidth and kernel launch overhead.

MLX's Metal backend implements GPU computation for Apple Silicon through Metal command encoding, JIT compilation of kernels, and device management. The backend translates primitives into Metal Shading Language (MSL) kernels, compiles them at runtime, and encodes commands into Metal command buffers for GPU execution. Device management abstracts Metal's GPU/CPU memory hierarchy, and stream abstraction enables asynchronous command submission and synchronization.

MLX's CUDA backend enables GPU computation on NVIDIA hardware through CUDA graph management, memory synchronization, and device abstraction. The backend manages CUDA device contexts, allocates GPU memory, and uses CUDA graphs to batch kernel launches for reduced overhead. Device management handles stream synchronization and memory pooling to optimize allocation patterns.

MLX provides a NumPy-like API for array operations (mlx.core) with familiar functions like reshape, transpose, matmul, conv2d, and element-wise operations. The API is implemented through Python bindings (Nanobind) that wrap C++ operations, enabling users familiar with NumPy to use MLX with minimal learning curve. Operations return lazy arrays that build computation graphs rather than executing immediately.

MLX provides mlx.nn, a neural network module system where layers inherit from a base Module class that manages parameters and submodules. The Module system tracks trainable parameters, enables parameter sharing, and provides methods for parameter initialization and state management. Layers (Linear, Conv2d, BatchNorm, etc.) are implemented as Modules that compose primitive operations, and the system integrates with autodiff for gradient computation during training.

MLX implements quantization (mlx.core.quantize) with multiple modes (int4, int8, float16) and backend-specific implementations for Metal and CUDA. Quantization reduces model size and inference latency by converting weights to lower precision, with dequantization happening on-the-fly during computation. The framework provides both quantization APIs for converting models and quantized operations that handle dequantization transparently.