Topology-Aware Distributed LLM Inference

Overview

Distributed LLM inference splits a model across multiple devices so that models too large for one machine can still run locally or in a cluster. The central challenge is placement: which device should hold which weights, which parallelism strategy should be used, and how much communication overhead will each generated token pay. Topology-aware inference systems answer those questions using a live view of device memory, compute capability, and link bandwidth.

Exo is a useful case study because it targets heterogeneous local clusters rather than only homogeneous datacenter GPUs. Its README describes automatic device discovery, topology-aware auto parallel placement, tensor parallel sharding, MLX distributed communication, and API compatibility with common chat protocols. The broader pattern applies anywhere an inference runtime must turn several uneven machines into one serving target.

The mechanism is different from ordinary request load balancing. A load balancer sends independent requests to different replicas. Distributed inference for an oversized model sends one request through multiple workers that jointly own the model state. That means token latency depends on both compute time and inter-device activation exchange, so placement must optimize memory fit and communication cost at the same time.

Core Concepts

Device Topology

A topology is a graph of devices and links. Nodes carry resource attributes such as memory, accelerator type, supported kernels, and current load. Edges carry link attributes such as latency, bandwidth, and whether the link supports a low-latency transport. A topology-aware planner uses this graph to avoid placing chatty shards across slow links.

Pipeline Parallelism

Pipeline parallelism assigns contiguous layer ranges to different devices. Device A computes early layers, sends activations to device B, and so on. It reduces per-device memory requirements and has lower communication volume than tensor parallelism, but a single request must still traverse the pipeline in order.

Tensor Parallelism

Tensor parallelism shards individual weight matrices across devices. Multiple devices compute parts of the same layer and then exchange partial results, often with all-reduce style communication. This can improve latency when links are fast enough, but the communication pattern is more demanding than pipeline-only placement.

How It Works

Discover Devices and Build the Topology

The runtime starts by discovering peers and collecting resource data. In a local cluster, peers may appear and disappear as laptops sleep, cables change, or processes restart. The control plane must refresh topology state and invalidate placements that no longer fit. Exo exposes cluster state and placement previews through HTTP endpoints so users can inspect possible instance layouts before creating a model instance.

{
  "nodes": {
    "mac-studio-a": {
      "memory_bytes": 549755813888,
      "accelerator": "mlx-metal",
      "links": {"mac-studio-b": {"latency_us": 18, "bandwidth_gbps": 80}}
    },
    "mac-studio-b": {
      "memory_bytes": 549755813888,
      "accelerator": "mlx-metal",
      "links": {"mac-studio-a": {"latency_us": 18, "bandwidth_gbps": 80}}
    }
  }
}

Preview Candidate Placements

Before loading a model, the planner enumerates valid placements. Each candidate records the model ID, sharding strategy, per-node memory delta, and any placement error. This preview step is useful because loading hundreds of gigabytes of weights is expensive; rejecting an invalid plan after loading would waste time and memory bandwidth.

curl "http://localhost:52415/instance/previews?model_id=llama-3.2-1b"

{
  "previews": [
    {
      "model_id": "mlx-community/Llama-3.2-1B-Instruct-4bit",
      "sharding": "Pipeline",
      "instance_meta": "MlxRing",
      "memory_delta_by_node": {"local": 729808896},
      "error": null
    }
  ]
}

Create the Instance and Route Requests

After choosing a placement, the runtime loads the assigned shards and exposes a compatible API. A chat-completions request looks ordinary to the client, but internally the request flows through the chosen distributed execution plan. The serving layer should hide placement details from clients while still exposing operational state for debugging and benchmarking.

Practical Applications

Local Clusters for Oversized Models

Personal or lab clusters can run models that exceed the memory of any single workstation. This is most compelling when privacy, offline operation, or hardware reuse matters more than the lowest possible latency. A cluster of several high-memory local machines can host a larger quantized model while preserving a familiar OpenAI-compatible API surface.

Heterogeneous Edge Inference

Edge deployments often mix devices with different memory sizes, accelerators, and network links. Topology-aware placement lets stronger devices hold larger shards, while weaker devices contribute only when the communication cost is justified. The same idea applies to offices, labs, production floors, and offline field deployments.

Benchmarking Parallelism Choices

Distributed inference systems need empirical validation because the best placement depends on prompt length, generation length, cache size, model architecture, and interconnect. Benchmark tools should report prompt tokens per second, generated tokens per second, peak memory, placement metadata, and repeated-run variance.

Implementation Patterns

Represent Placement as Data

Make placements serializable API objects rather than hidden runtime decisions. This allows preview, approval, replay, and benchmarking. A serialized placement also gives operators a stable handle for deleting, recreating, or comparing model instances.

type Placement struct {
    ModelID       string
    Strategy      string // "pipeline", "tensor", or "hybrid"
    Runtime       string // "mlx-ring", "jaccl", "cuda", etc.
    Shards        []ShardAssignment
    EstimatedCost CostEstimate
}

type ShardAssignment struct {
    NodeID      string
    LayerStart  int
    LayerEnd    int
    MemoryBytes int64
}

Separate Control Plane from Data Plane

The control plane discovers peers, calculates placements, creates instances, and exposes state. The data plane performs token generation and communication. Keeping these paths separate prevents a slow dashboard or placement preview from interfering with active inference.

Choose Parallelism from Link Quality

Use pipeline parallelism when cross-node links are slower or when memory fit is the primary goal. Use tensor parallelism when links have high bandwidth and low latency, because tensor sharding exchanges partial results inside layers. Hybrid plans are useful when a cluster has fast links within a node group and slower links between groups.

Performance Considerations

• Interconnect cost: Tensor parallelism can be bottlenecked by all-reduce traffic. Fast links such as NVLink, InfiniBand, or low-latency local transports are often required for good scaling.
• KV cache headroom: Fitting weights is not enough. Long prompts and high concurrency require cache memory, and adding devices can increase usable cache capacity.
• Pipeline bubbles: Pipeline parallelism can leave stages idle unless requests or micro-batches keep the pipeline full. Single-stream interactive chat may scale worse than batch serving.
• Heterogeneous devices: A slow device in the middle of a pipeline can dominate latency. Weighted placement should account for both memory and compute throughput.
• Startup overhead: Loading large sharded models can take long enough that previewing placements and caching model files become operational requirements.
• Transport stability: Local clusters are sensitive to cable quality, operating-system versions, peer discovery, and sleep behavior. Treat topology as dynamic state, not as a static config file.

Common Pitfalls

• Assuming more devices always means faster tokens: Extra devices add communication overhead. Distributed inference is often justified by memory capacity before latency improvement.
• Ignoring prompt versus decode phases: Prefill and decode stress different parts of the system. A placement that looks good for short prompts may underperform on long-context workloads.
• Hiding placement decisions: Operators need to know where shards were placed and why. Without this, performance regressions are difficult to explain.
• Overloading one compatibility endpoint: Supporting OpenAI, Claude-style, and Ollama-style APIs is convenient, but each protocol has different streaming and error semantics. Normalize internally but preserve client expectations.
• Trusting remote model code by default: Custom model loading can require executable code from model repositories. Keep remote-code execution disabled unless the source is explicitly trusted.