How to Use Apple Silicon Unified Memory for LLM Inference

Introduction to Apple Silicon Unified Memory for LLM Inference

Apple Silicon chips (M1, M2, M3, M4 series) feature a groundbreaking architecture called Unified Memory Architecture (UMA). Unlike traditional systems where CPU and GPU have separate memory pools connected via a bus, Apple Silicon uses a single pool of high-bandwidth memory shared between the CPU, GPU, and Neural Engine. This design has profound implications for running Large Language Models (LLMs) locally.

In a traditional discrete GPU setup, you must copy model weights from system RAM (CPU memory) to GPU VRAM before inference. This copy operation is expensive and limits you to models that fit in the GPU's VRAM. With Apple Silicon, the GPU can directly access the same memory the CPU uses, eliminating the copy bottleneck and allowing you to load models up to the size of your total system RAM.

Key Specifications Across Apple Silicon Generations

• M1/M2/M3/M4 (base): Up to 24GB unified memory, ~100 GB/s bandwidth
• M1/M2/M3 Pro: Up to 36GB unified memory, ~200 GB/s bandwidth
• M1/M2/M3/M4 Max: Up to 128GB unified memory, ~400 GB/s bandwidth
• M2/M3 Ultra: Up to 192GB unified memory, ~800 GB/s bandwidth

The memory bandwidth is particularly important for LLM inference, which is heavily memory-bound. Each token generation requires reading the entire model weights from memory, so higher bandwidth directly translates to faster token generation.

Why Unified Memory Matters for LLM Inference

The Memory Wall Problem

LLM inference is fundamentally a memory-bandwidth-bound problem. For each token generated, the system must read all model parameters from memory. A 7B parameter model in 4-bit quantization is roughly 3.5GB, meaning every single token requires reading 3.5GB of data. At 100 GB/s bandwidth, this limits you to about 28 tokens per second — and that is the theoretical maximum before accounting for any overhead.

This is why Apple Silicon's high memory bandwidth is so valuable. An M2 Ultra with 800 GB/s can theoretically generate tokens from a 70B model (quantized to ~35GB) at over 20 tokens per second, which is remarkable for a desktop-class machine.

Eliminating the CPU-to-GPU Copy

On traditional systems with NVIDIA or AMD GPUs, loading a model requires transferring weights from system RAM to GPU VRAM over PCIe. For a 13GB model, this transfer can take several seconds. With Unified Memory, the model weights are loaded directly into shared memory once, and both the CPU and GPU can access them immediately. This dramatically reduces load times and simplifies the inference pipeline.

Running Larger Models Than VRAM Allows

A consumer NVIDIA RTX 4090 has 24GB of VRAM. To run a 70B parameter model, you would need multiple GPUs or aggressive quantization. An Apple Silicon Mac with 128GB or 192GB of unified memory can load that same model entirely in memory with room to spare, all on a single chip. This democratizes access to large model inference for developers who do not have access to multi-GPU clusters.

How to Use Apple Silicon for LLM Inference

There are several frameworks optimized for Apple Silicon. The two most popular are MLX (Apple's own machine learning framework) and llama.cpp (a C++ inference engine with Metal support). We will cover both approaches.

Approach 1: Using MLX and mlx-lm

MLX is an array framework for machine learning on Apple Silicon, developed by Apple's machine learning research team. The mlx-lm package provides a high-level API for running LLMs optimized for Unified Memory.

First, install the required packages:

pip install mlx-lm

Here is a complete example of loading and running inference with a model using MLX:

from mlx_lm import load, generate

# Load a model - this downloads from HuggingFace and converts to MLX format
# The 4-bit quantized version fits easily in 16GB unified memory
model, tokenizer = load(
    "mlx-community/Mistral-7B-Instruct-v0.3-4bit"
)

# Create a prompt using the chat template
messages = [
    {"role": "user", "content": "Explain how unified memory works on Apple Silicon."}
]
prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

# Generate a response
response = generate(
    model,
    tokenizer,
    prompt=prompt,
    max_tokens=512,
    verbose=True  # Prints token generation speed
)

print(response)

MLX automatically uses the Metal GPU for computation and leverages Unified Memory so no weight copying occurs. The verbose=True flag will print the tokens-per-second rate, allowing you to measure performance.

Streaming Responses with MLX

For interactive applications, you typically want to stream tokens as they are generated. MLX provides a streaming API for this:

from mlx_lm import load, stream_generate

model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-4bit")

messages = [
    {"role": "system", "content": "You are a helpful coding assistant."},
    {"role": "user", "content": "Write a Python function to compute fibonacci numbers."}
]

prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

print("Response: ", end="", flush=True)
for token in stream_generate(model, tokenizer, prompt, max_tokens=256):
    print(token.text, end="", flush=True)
print()  # Newline after completion

Converting and Quantizing Models for MLX

While the mlx-community HuggingFace organization provides many pre-quantized models, you may want to convert and quantize your own models. MLX includes a CLI tool for this:

# Convert and quantize a HuggingFace model to 4-bit MLX format
python -m mlx_lm.convert \
    --hf-path mistralai/Mistral-7B-Instruct-v0.3 \
    --mlx-path ./my-mistral-4bit \
    --quantize \
    --q-bits 4 \
    --q-group-size 64

# Then load your local converted model
python -m mlx_lm.generate \
    --model ./my-mistral-4bit \
    --prompt "Hello, how are you?" \
    --max-tokens 100 \
    --verbose

Approach 2: Using llama.cpp with Metal

llama.cpp is another excellent option that supports Apple's Metal framework for GPU acceleration. It is particularly useful if you need broader model format support or want a lightweight C++ binary.

Build llama.cpp with Metal support:

# Clone the repository
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp

# Build with Metal support (enabled by default on macOS)
make

# Verify Metal is available
./llama-cli --version

Run inference with a GGUF model file:

# Download a GGUF model (example: Llama 3.2 3B Q4_K_M)
# from https://huggingface.co/models?library=gguf

# Run inference using the Metal GPU
./llama-cli \
    -m ./models/Llama-3.2-3B-Instruct-Q4_K_M.gguf \
    -p "Explain quantum computing in simple terms." \
    -n 256 \
    --gpu-layers 99 \
    --threads 4 \
    --color

The --gpu-layers 99 flag tells llama.cpp to offload all layers to the Metal GPU. Since Unified Memory means there is no VRAM limitation separate from system RAM, you can offload the entire model to the GPU.

Using llama.cpp Python Bindings

For integration into Python applications, you can use the llama-cpp-python package with Metal support:

# Install with Metal support
CMAKE_ARGS="-DGGML_METAL=on" pip install llama-cpp-python

# Python inference example
from llama_cpp import Llama

llm = Llama(
    model_path="./models/Llama-3.2-3B-Instruct-Q4_K_M.gguf",
    n_gpu_layers=-1,  # Offload all layers to Metal GPU
    n_ctx=4096,       # Context window size
    verbose=False
)

response = llm.create_chat_completion(
    messages=[
        {"role": "user", "content": "What are the advantages of unified memory?"}
    ],
    max_tokens=256,
    stream=False
)

print(response["choices"][0]["message"]["content"])

Approach 3: Using PyTorch with MPS Backend

PyTorch supports Apple Silicon through the MPS (Metal Performance Shaders) backend. While not as optimized as MLX for LLM workloads, it is useful if you already have PyTorch-based code:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Verify MPS is available
if not torch.backends.mps.is_available():
    raise RuntimeError("MPS is not available on this system")

device = torch.device("mps")

model_id = "meta-llama/Llama-3.2-1B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map=device
)

messages = [{"role": "user", "content": "Write a haiku about programming."}]
input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt"
).to(device)

with torch.no_grad():
    output = model.generate(
        input_ids,
        max_new_tokens=100,
        do_sample=True,
        temperature=0.7
    )

response = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True)
print(response)

Note that PyTorch with MPS is generally slower than MLX for LLM inference because MLX is specifically designed to exploit Unified Memory, while PyTorch's MPS backend has more overhead.

Best Practices for Apple Silicon LLM Inference

Choose the Right Quantization Level

Quantization is critical for efficient inference on Apple Silicon. The trade-off is between model quality and memory usage plus speed. Here are practical recommendations:

• 4-bit quantization (Q4): Best balance for most use cases. A 7B model uses ~3.5GB, a 70B model uses ~35GB.
• 8-bit quantization (Q8): Higher quality but double the memory. Use when you have ample RAM and need maximum quality.
• 3-bit quantization (Q3): Aggressive compression for very large models on machines with limited memory. Quality degradation is noticeable.

For group size, 64 is the standard and works well. Smaller group sizes (like 32) slightly improve quality but increase memory usage.

Manage Memory Efficiently

Even with Unified Memory, you have a finite pool. macOS reserves a portion of RAM for the system, and the GPU can only use a subset of total unified memory. As a rule of thumb, the GPU can access approximately 70-80% of total system RAM. On a 32GB Mac, plan for about 24-25GB available for model weights and KV cache.

To check available memory before loading a model:

import subprocess
import psutil

# Check total and available system memory
mem = psutil.virtual_memory()
print(f"Total RAM: {mem.total / (1024**3):.1f} GB")
print(f"Available RAM: {mem.available / (1024**3):.1f} GB")

# Check Metal GPU memory limits via system_profiler
result = subprocess.run(
    ["system_profiler", "SPHardwareDataType"],
    capture_output=True, text=True
)
print(result.stdout)

Optimize the KV Cache

The KV cache grows with context length and can consume significant memory. For a 7B model with a 4096 token context, the KV cache can use 1-2GB depending on precision. MLX supports quantized KV cache to reduce this:

from mlx_lm import load, generate

model, tokenizer = load(
    "mlx-community/Mistral-7B-Instruct-v0.3-4bit",
    tokenizer_config={"trust_remote_code": True}
)

# Use quantized KV cache by specifying kv_bits
response = generate(
    model,
    tokenizer,
    prompt="Summarize the following text: ...",
    max_tokens=512,
    kv_bits=8,  # Quantize KV cache to 8-bit to save memory
    verbose=True
)

Batch Inference for Throughput

If you need to process many prompts, batch inference is more efficient than sequential calls. MLX supports batching natively:

from mlx_lm import load, generate
import time

model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-4bit")

prompts = [
    "What is machine learning?",
    "Explain neural networks.",
    "How does backpropagation work?",
    "What is gradient descent?"
]

# Format all prompts
formatted = [
    tokenizer.apply_chat_template(
        [{"role": "user", "content": p}],
        tokenize=False,
        add_generation_prompt=True
    )
    for p in prompts
]

# Batch generate
start = time.time()
responses = generate(
    model,
    tokenizer,
    prompt=formatted,
    max_tokens=200,
    verbose=True
)
elapsed = time.time() - start

for i, resp in enumerate(responses):
    print(f"\n--- Prompt {i+1} ---")
    print(resp)

print(f"\nTotal time: {elapsed:.2f}s for {len(prompts)} prompts")

Choose the Right Framework for Your Use Case

• MLX / mlx-lm: Best for Python-native development, research, and when you want maximum Apple Silicon optimization. First-class support for Unified Memory.
• llama.cpp: Best for production deployments, server applications, and when you need the widest model format support (GGUF). Excellent Metal integration.
• PyTorch + MPS: Best when you already have PyTorch codebases or need custom training/fine-tuning pipelines. Less optimized for pure inference.
• Ollama: Best for quick local setup and API-compatible serving. Under the hood, it uses llama.cpp with Metal.

Monitor Performance

Always measure your actual tokens-per-second to understand performance. You can use the built-in MLX verbose output or measure manually:

from mlx_lm import load, generate
import time

model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")

prompt = "Write a detailed essay about artificial intelligence."

start = time.time()
response = generate(
    model,
    tokenizer,
    prompt=prompt,
    max_tokens=500,
    verbose=False
)
elapsed = time.time() - start

tokens_generated = len(tokenizer.encode(response))
tps = tokens_generated / elapsed

print(f"Generated {tokens_generated} tokens in {elapsed:.2f}s")
print(f"Tokens per second: {tps:.1f}")
print(f"\nResponse:\n{response}")

Keep Models on Disk in the Right Format

Downloading and converting models every time is wasteful. Store converted MLX models or downloaded GGUF files in a dedicated directory. For MLX, models are cached in ~/.cache/huggingface/hub/ by default. You can also save converted models locally and load them by path:

from mlx_lm import load, save

# Load from HuggingFace
model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-4bit")

# Save locally for future use
save("my-local-models/llama-3.2-3b-4bit", model, tokenizer)

# Later, load from local path
model, tokenizer = load("my-local-models/llama-3.2-3b-4bit")

Conclusion

Apple Silicon's Unified Memory Architecture represents a paradigm shift for local LLM inference. By eliminating the CPU-to-GPU memory copy bottleneck and allowing the GPU to access the full system memory pool, Apple Silicon makes it practical to run large language models on consumer and workstation-class machines without expensive multi-GPU setups. Whether you choose MLX for its native optimization, llama.cpp for its production readiness, or PyTorch for compatibility with existing code, the key to success lies in understanding the memory-bandwidth-bound nature of LLM inference and making informed choices about quantization, KV cache management, and framework selection. With the practical examples and best practices covered in this tutorial, you are well-equipped to leverage Apple Silicon for efficient, high-performance LLM inference in your own projects.