Distributed Inference Overview¶

Table of Contents¶

• Why Distribute?
• Strategy Comparison
• When to Use Each Strategy
• Common Infrastructure
• Cluster Configuration

Why Distribute?¶

Modern LLMs can require tens or hundreds of gigabytes of memory just for model weights alone. A 70B parameter model in 16-bit precision occupies roughly 140 GB, which exceeds the capacity of any single consumer or workstation GPU. Even when the model fits, running multiple concurrent requests can exhaust device memory for KV caches.

Distribution solves two distinct problems:

1. Memory capacity: the model is too large for one GPU. Solution: split the model so each GPU holds only a portion.

2. Throughput: a single GPU cannot keep up with the desired request rate. Solution: run multiple copies of the model or pipeline stages in parallel.

mlxcel provides three complementary strategies that can be combined depending on the deployment topology.

Strategy Comparison¶

graph TB
    subgraph TP["Tensor Parallelism (TP)"]
        direction LR
        TP_IN[Input Tokens] --> TP_R0[Rank 0<br/>W shard 0]
        TP_IN --> TP_R1[Rank 1<br/>W shard 1]
        TP_R0 -->|all-reduce| TP_OUT[Output]
        TP_R1 -->|all-reduce| TP_OUT
    end

    subgraph PP["Pipeline Parallelism (PP)"]
        direction LR
        PP_S0[Stage 0<br/>Layers 0-15] -->|activations| PP_S1[Stage 1<br/>Layers 16-31]
        PP_S1 -->|activations| PP_S2[Stage 2<br/>Layers 32-47]
        PP_S2 --> PP_OUT[Logits]
    end

    subgraph DI["Disaggregated Inference (DI)"]
        direction LR
        DI_IN[Request] --> DI_PF[Prefill Node<br/>Compute KV cache]
        DI_PF -->|transfer KV| DI_DC[Decode Node<br/>Autoregressive decode]
        DI_DC --> DI_OUT[Tokens]
    end

When to Use Each Strategy¶

Use Tensor Parallelism when:

• You have multiple GPUs on the same host (or connected via NVLink/Thunderbolt).
• The model does not fit on a single GPU but fits across 2-8 GPUs.
• You need low per-request latency (TP all-reduce is fast on shared-memory busses).
• The model is a standard Transformer, MoE, or SSM architecture (all are supported with architecture-aware shard plans).

Use Pipeline Parallelism when:

• You have multiple hosts and the model is too large for any single host.
• Per-request latency can tolerate pipeline bubble overhead.
• Requests arrive in predictable batch sizes (micro-batching amortizes pipeline bubbles effectively).
• You want to maximize total throughput over a set of heterogeneous nodes.

Use Disaggregated Inference when:

• You serve workloads with highly variable prompt lengths (e.g., RAG, code completion with long context).
• Prefill nodes and decode nodes have different hardware profiles (compute-dense GPUs for prefill, memory-dense GPUs for decode).
• You want to scale prefill capacity and decode capacity independently without changing the other.
• Long prompts should not delay short decode sessions that are already in flight.

Common Infrastructure¶

All three strategies share a common infrastructure layer defined in src/distributed/.

Node Discovery¶

Each node advertises itself with a unique ID and network address. The cluster configuration (TOML or CLI flags) provides the initial peer list.

Heartbeat and Failure Detection¶

A background heartbeat loop (src/distributed/heartbeat.rs) sends periodic pings to all known peers. The failure detector (src/distributed/failure_detector.rs) marks a node as unavailable after a configurable number of missed heartbeats, triggering re-routing in the request router.

Connection Pooling¶

src/distributed/connection_pool.rs maintains persistent TCP connections to peer nodes. Connections are reused across requests to avoid TCP handshake overhead on the critical path.

Transport Layer¶

The transport abstraction (src/distributed/transport.rs) supports:

• TCP: Default transport, works across all network topologies.
• Thunderbolt: High-bandwidth transport for Apple Silicon devices connected via Thunderbolt; see src/distributed/thunderbolt_transport.rs.
• RDMA: Optional accelerated path negotiated per peer at connect time, with a TCP fallback; see src/distributed/rdma_transport.rs and Pipeline Parallelism.

graph LR
    A[Node A] -->|TCP / Thunderbolt| B[Node B]
    B --> C[Node C]
    C --> A
    style A fill:#4a90d9
    style B fill:#4a90d9
    style C fill:#4a90d9

Cluster Secret, Handshake, and Frame Protection¶

Pipeline roles and the disaggregated prefill, decode, and router roles refuse to start without a pre-shared cluster secret (--cluster-secret-file, MLXCEL_CLUSTER_SECRET_FILE, MLXCEL_CLUSTER_SECRET, or [cluster] secret_file in the cluster TOML; at least 16 bytes). Single-node serving (hybrid) does not read it. The secret is used in three places:

• Handshake (src/distributed/cluster_handshake.rs, protocol version 5). Each TCP connection starts with a mutual challenge-response: both ends prove possession of the secret with HMAC-SHA256 over the full transcript, so a peer with the wrong secret is rejected on whichever side it is. The accept side binds the client's advertised identity to the TCP source IP. The server's verdict (accept, or at capacity) carries its own HMAC tag, so a party on the path cannot change one verdict into the other.
• Frame protection (src/distributed/frame_protection.rs). Per-connection directional keys are derived from the secret and the handshake nonces with HKDF-SHA256. Every frame after the handshake is encrypted and authenticated with AES-256-GCM, the default suite, and carries an implicit sequence number, so a replayed, reordered, dropped, or injected frame fails authentication and closes the connection. Both ends must use the same suite; the handshake refuses a mismatch.
• Discovery beacons (src/distributed/discovery_beacon.rs) are signed with the same secret.

Connections that have not finished the handshake draw from a separate pending pool with a per-source limit (src/distributed/pending_handshakes.rs), so a flood of idle sockets cannot use up the slots of established connections.

Tensor Protocol¶

src/distributed/tensor_protocol.rs defines the wire format for transferring tensors between nodes. Supports the following data types:

For KV cache transfer in disaggregated inference, optional LZ4 compression and on-the-fly int8/int4 quantization are available to reduce bandwidth requirements.

Cluster Configuration¶

A cluster is configured via a TOML file passed with --cluster-config, or via CLI flags (--node-id, --node-role, --peers) for simpler topologies.

TOML Configuration¶

[cluster]
name = "my-cluster"
tensor_parallel_size = 2
pipeline_parallel_size = 1

[[nodes]]
id = "node-0"
address = "192.168.1.10:8080"
role = "tensor_parallel_rank"
rank = 0

[nodes.resources]
memory_bytes = 68719476736   # 64 GiB
compute_units = 0            # 0 = auto-detect

[[nodes]]
id = "node-1"
address = "192.168.1.11:8080"
role = "tensor_parallel_rank"
rank = 1

[nodes.resources]
memory_bytes = 68719476736
compute_units = 0

Node Roles¶

--node-role router additionally selects the model-free HTTP front of a disaggregated deployment. It is a CLI role, not a TOML node role; see Disaggregated Inference.

CLI Configuration (simple topologies)¶

# Node 0 (rank 0)
mlxcel-server --model ./models/llama-70b \
  --node-id node-0 \
  --node-role tensor_parallel_rank \
  --peers 192.168.1.11:8080

# Node 1 (rank 1)
mlxcel-server --model ./models/llama-70b \
  --node-id node-1 \
  --node-role tensor_parallel_rank \
  --peers 192.168.1.10:8080

Validation¶

The cluster configuration is validated at startup:

• Node IDs must be unique and non-empty.
• Addresses must be unique.
• Node IDs and cluster name must not exceed 128 characters and must not contain control characters.
• At least one node must be present.