How to Run Giant LLMs Locally with MLX and RDMA

Peter Bubenik · Apple ML · · Source
Image for Local and Distributed Large LLM Inference with MLX and RDMA on Mac Studio

After studying this material, you should be able to:

  1. Explain how MLX enables LLM inference on Apple Silicon hardware
  2. Describe the difference between pipeline and tensor parallelization
  3. Understand why certain metrics (TTFT vs. decode speed) are bound by different hardware resources
  4. Apply MLX commands to benchmark and serve LLMs across single or multiple machines
  5. Evaluate tradeoffs between parallelization strategies for different model sizes and contexts

Step-by-Step Study Material


Step 1: Foundation — What Is MLX and Why Does It Matter?

What is MLX?

MLX is an open-source array framework built specifically for Apple Silicon. Think of it like PyTorch or NumPy, but optimized for Mac hardware.

Why Apple Silicon is Special for AI

Apple Silicon uses a Unified Memory Architecture (UMA):

Traditional Hardware:          Apple Silicon:
┌─────────┐  ┌─────────┐      ┌──────────────────────┐
│   CPU   │  │   GPU   │      │   Unified Memory      │
│  RAM    │  │  VRAM   │      │  CPU + GPU share RAM  │
└─────────┘  └─────────┘      └──────────────────────┘
Data must copy between them    No copying needed = faster

Key Hardware Specs to Know

ChipMemory BandwidthMax RAM
M3 Ultra819 GB/s512 GB
M5 Ultra1.2 TB/s512 GB

Why this matters: More bandwidth = faster token generation


Step 2: Understanding LLM Inference Metrics

Before benchmarking makes sense, you need to understand what is being measured.

Metric 1: Time to First Token (TTFT)

  • Definition: How long before the model outputs its first word
  • Bottleneck: Compute-bound (processing the input/context)
  • Analogy: How long a chef takes to read your order before cooking

Metric 2: Decode Throughput (tokens/second)

  • Definition: How fast the model generates output tokens
  • Bottleneck: Memory-bandwidth-bound (loading model weights repeatedly)
  • Analogy: How fast the chef actually cooks each dish

Why the Bottleneck Difference Matters

TTFT (Prefill Phase):
Input tokens → [Matrix Multiplications] → First output token
                ↑
         Needs fast COMPUTE
         → M5 Neural Accelerators help here (up to 4.6× faster)

Decode Phase:
Each new token → [Load weights from RAM] → Next token
                  ↑
         Needs fast MEMORY BANDWIDTH
         → M5's 1.2 TB/s helps here (1.6× faster)

Step 3: Model Types — Dense vs. Mixture of Experts (MoE)

Dense Models

  • All parameters are active during every inference
  • Examples: Mistral Medium 3.5 (128B), Qwen 3.8 (27B)
  • More predictable memory usage

Mixture of Experts (MoE)

  • Only a fraction of parameters are active per token
  • Example: Muse Spark Open (816B total, but only 42B active)
  • Reduces compute at inference
Dense Model (128B):
Token → [ALL 128B parameters process it] → Output

MoE Model (816B total, 42B active):
Token → [Router] → [Expert 1] [Expert 7] [Expert 23]... → Output
                    Only ~42B parameters used per token

Precision Formats

FormatBits per ParameterUse Case
BF1616Standard training/inference
FP88Efficient inference
NVFP44Aggressive compression

Key insight: Muse Spark Open needs ~1.6TB in BF16, but only ~493GB in NVFP4 — making it possible to run on Mac hardware


Step 4: Single Machine Inference

What Fits on One Mac Studio (512GB)?

Qwen 3.8 (27B):    ████░░░░░░░░░░░░  ~68GB  ✓ Fits easily
Mistral (128B):    ████████░░░░░░░░  ~158GB ✓ Fits
Muse Spark (816B): ████████████████  ~493GB+ ✗ Does NOT fit

Running a Benchmark (Single Machine)

mlx_lm.benchmark --model Qwen/Qwen3.8-27B -p 8192 -g 128 -b 1

Breaking down the flags:

  • --model → which model to load
  • -p 8192 → 8,192 tokens of input context
  • -g 128 → generate 128 output tokens
  • -b 1 → batch size of 1

Performance Results (Single Machine)

ModelM3 Ultra TTFTM5 Ultra TTFTSpeedup
Qwen 3.818.2s4.1s4.4×
Mistral 128B95.9s24.6s3.9×

Step 5: Scaling Across Multiple Machines

Why You Need Multiple Machines

When a model is too large for one machine, you must distribute it.

The Communication Technology: RDMA over Thunderbolt

  • RDMA = Remote Direct Memory Access
  • Allows one machine to read/write another machine's memory directly, bypassing the CPU
  • Thunderbolt 5 provides the physical connection
  • MLX uses a communication layer called jaccl (pronounced "jackal") built on this
Mac Studio 1 ←──Thunderbolt 5──→ Mac Studio 2
              (RDMA / jaccl)
              Low latency, high bandwidth

Step 6: Pipeline Parallelization

Concept

Split the model by depth (layers):

Model Layers: [1][2][3][4][5][6][7][8]
                    ↓
Mac 1: [1][2][3][4]  →  sends output  →  Mac 2: [5][6][7][8]

How Prefill Pipelining Works

For large contexts, MLX breaks input into chunks and overlaps computation:

Time →
Mac 1: [Chunk 1, layers 1-4] [Chunk 2, layers 1-4] ...
Mac 2:                        [Chunk 1, layers 5-8] [Chunk 2, layers 5-8]
                              ↑ Overlap! Both machines work simultaneously

Running Pipeline Parallelization

mlx.launch --hostfile hostfile.json --backend jaccl \
  --env MLX_METAL_FAST_SYNCH=1 -- \
  mlx_lm.benchmark --model mistralai/Mistral-Medium-3.5-128B \
  -p 8192 -g 128 -b 1 --prefill-step-size 8192 --pipeline

Key flags:

  • --backend jaccl → use RDMA Thunderbolt
  • MLX_METAL_FAST_SYNCH=1 → faster CPU-GPU sync
  • --pipeline → enable pipeline parallelization
  • --prefill-step-size 8192 → chunk size for overlapping

Pipeline Parallelization Results

MetricResult
Decode speedSimilar to single machine
TTFT at large contextUp to 1.8× faster (due to overlap)

Limitation: Layers still run sequentially, so decode doesn't improve much


Step 7: Tensor Parallelization

Concept

Split the model by width (each layer distributed across machines):

Pipeline (depth split):          Tensor (width split):
Mac 1: [Layer 1-4]               Mac 1: [Half of EVERY layer]
Mac 2: [Layer 5-8]               Mac 2: [Half of EVERY layer]
Sequential                       Parallel per layer

Visual Comparison

Pipeline Parallelization:
Mac1 → [L1→L2→L3→L4] → Mac2 → [L5→L6→L7→L8] → Output
                         ↑ waits for Mac1

Tensor Parallelization:
Mac1 → [L1a→L2a→L3a...] ↘
                          → Combine → Output (each layer)
Mac2 → [L1b→L2b→L3b...] ↗
Both work simultaneously per layer!

Running Tensor Parallelization

mlx.launch --hostfile hostfile.json --backend jaccl \
  --env MLX_METAL_FAST_SYNCH=1 -- \
  mlx_lm.benchmark --model meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
  -p 8192 -g 128 -b 1

Note: No --pipeline flag needed — tensor parallelization is the default when launching across multiple nodes

Tensor Parallelization Results

Model1 Machine2 Machines4 Machines
Muse Spark Open decodebaseline24.99 t/s36.73 t/s
Mistral 128B decode7.70 t/s13.69 t/s23.37 t/s
Qwen 3.8 decode19.69 t/s33.27 t/s50.89 t/s

Step 8: Choosing the Right Strategy

Decision Framework

Is the model too large for one machine?
        │
        ├── NO → Use single machine
        │
        └── YES → How many machines do you have?
                        │
                        ├── 2 machines, large context → Pipeline may be better for TTFT
                        │
                        └── 4+ machines → Tensor parallelization wins

Tradeoffs Summary

StrategyComplexityDecode SpeedTTFTCommunication
Single MachineLowBaselineBaselineNone
PipelineLow≈ SameBetter at large contextLow (only between layer groups)
TensorHighMuch betterBetter (scales with machines)High (every layer, every token)

Key insight: Tensor parallelization requires very low latency communication — this is why RDMA over Thunderbolt is critical


Step 9: Serving LLMs Locally for Agentic AI

Starting a Local LLM Server

mlx.launch --hostfile hostfile.json --backend jaccl \
  --env MLX_METAL_FAST_SYNCH=1 -- \
  mlx_lm.server --model meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
  --port 8080

Connecting an AI Coding Tool (OpenCode Example)

{
  "model": "meta-models/Muse-Spark-Open-816B-A42B-NVFP4",
  "provider": {
    "mlx": {
      "options": {
        "baseURL": "http://127.0.0.1:8080/v1"
      }
    }
  }
}

The server exposes an OpenAI-compatible API, so any tool that supports OpenAI can use your local model.


Step 10: Putting It All Together

The Full Picture

Hardware Layer:
Mac Studio M5 Ultra
├── 512GB Unified Memory (CPU + GPU share)
├── 1.2 TB/s Memory Bandwidth
├── Neural Accelerators (fast matrix multiply)
└── Thunderbolt 5 (RDMA to other Macs)

Software Layer (MLX):
├── mlx_lm.benchmark  → measure performance
├── mlx_lm.server     → serve models via API
├── mlx.launch        → distribute across machines
└── jaccl backend     → RDMA communication

Parallelization:
├── Pipeline → split by layers (simple, good for very large models)
└── Tensor   → split by width (complex, better decode scaling)

Key Numbers to Remember

AchievementValue
Muse Spark Open on 4× M5 Ultra36+ tokens/second
TTFT speedup M5 vs M3Up to 4.6×
Decode speedup M5 vs M31.6×
Peak power consumption~300W (M5 Ultra)
Largest model run locally816B parameters

Quick Knowledge Check

Test yourself with these questions:

  1. Why is TTFT compute-bound but decode is memory-bandwidth-bound?

    TTFT processes all input tokens at once (heavy computation). Decode generates one token at a time, requiring model weights to be loaded from memory repeatedly.

  2. Why does pipeline parallelization not improve decode speed much?

    Layers still run sequentially — one machine must wait for the other's output before proceeding.

  3. Why does tensor parallelization require low-latency communication?

    Every layer, for every token, requires broadcasting inputs and aggregating outputs across machines. High latency would negate the compute gains.

  4. What makes MoE models like Muse Spark Open memory-efficient at inference?

    Only a fraction of experts (and thus parameters) are active per token, reducing compute even though total parameter count is huge.

  5. What is the role of NVFP4 precision for Muse Spark Open?

    It reduces memory from ~1.6TB (BF16) to ~493GB, making it possible to fit across a cluster of Mac Studios.

More to study