Long prompts and serving many agents

Long prompts are slow to start

From Apple's post, for Muse Spark Open on four M5 Ultra Mac Studios with tensor parallel:

• 8,192-token context: 4.8 s to the first token.
• 65,536-token context: 40.6 s with two Thunderbolt cables per pair, and about 50.5 s with one.

Agent tools often send large contexts. So a coding agent on this cluster can wait about 40 seconds before the first word, on the best hardware Apple tested. Older Macs are slower: for Mistral Medium 3.5 at 8,192 tokens, one M3 Ultra took 95.9 s and one M5 Ultra took 24.6 s. Design for it: stream output, show that work is in progress, and keep the start of each prompt the same between calls so the prompt cache can reuse it (see Keep the prompt cache warm).

Tune the prefill chunk size

MLX processes a long prompt in chunks. mlx_lm.server uses --prefill-step-size 2048 by default (mlx-lm 0.32.2). Apple's pipeline benchmark used 8,192, and chunking is what lets pipeline overlap work across Macs. We did not measure the best value. Measure it on your cluster with the benchmark tool Apple used:

cd ~/mlx-cluster
for step in 2048 4096 8192; do
  uv run mlx.launch --hostfile jaccl.json --backend jaccl -- \
    ~/mlx-cluster/.venv/bin/python -m mlx_lm benchmark --model "$MODEL" \
    -p 65536 -g 128 -b 1 --prefill-step-size "$step"
done

Add --pipeline if you run pipeline. Write down the first-token time and peak memory for each value.

Run the server

uv run mlx.launch --hostfile jaccl.json --backend jaccl -- \
  ~/mlx-cluster/.venv/bin/python -m mlx_lm server \
  --model "$MODEL" --host 0.0.0.0 --port 8080 --prefill-step-size 8192

Facts from the mlx-lm source and docs that change how you run it:

• Only rank 0, the first host in the hostfile, runs the HTTP server. Other ranks wait for work. Point agents at rank 0.
• The default --host is 127.0.0.1, so other machines cannot reach it until you set --host.
• SERVER.md says the server is not recommended for production because it only has basic security checks. Keep it on a private network.
• In distributed mode, adapters and draft models are not supported.
• Batching: --decode-concurrency (default 32) and --prompt-concurrency (default 8). --kv-bits turns batching off, so requests run one at a time.
• We started a two-rank tensor parallel server on our Linux CPU box with the ring backend and got a correct chat answer. That only proves the command shape. It says nothing about Mac speed.

Apple's post shows how to point OpenCode at the server with an OpenAI-compatible provider and baseURL http://127.0.0.1:8080/v1. Use rank 0's address if the agent runs on another machine.

Measure what each agent waits

This script sends several streaming requests at once and prints the first-token time for each:

"""Send N agent requests at once and measure time to first token for each."""
import argparse
import json
import statistics
import threading
import time
import urllib.request


def one_request(url, prompt, max_tokens, results, i):
    body = json.dumps({
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "stream": True,
    }).encode()
    req = urllib.request.Request(url, body, {"Content-Type": "application/json"})
    start = time.perf_counter()
    first = None
    chunks = 0
    with urllib.request.urlopen(req, timeout=600) as resp:
        for raw in resp:
            line = raw.decode().strip()
            if not line.startswith("data:") or line == "data: [DONE]":
                continue
            delta = json.loads(line[5:])["choices"][0].get("delta", {})
            if delta.get("content") or delta.get("reasoning"):
                chunks += 1
                if first is None:
                    first = time.perf_counter() - start
    total = time.perf_counter() - start
    results[i] = (first, total, chunks)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--url", default="http://127.0.0.1:8080/v1/chat/completions")
    p.add_argument("--agents", type=int, default=4)
    p.add_argument("--prompt-words", type=int, default=2000)
    p.add_argument("--max-tokens", type=int, default=64)
    a = p.parse_args()

    filler = " ".join(["word"] * a.prompt_words)
    results = [None] * a.agents
    threads = [
        threading.Thread(target=one_request, args=(
            a.url, f"Agent {i}. Read this and reply with one short line.\n{filler}",
            a.max_tokens, results, i))
        for i in range(a.agents)
    ]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    ttfts = [r[0] for r in results if r and r[0] is not None]
    for i, r in enumerate(results):
        print(f"agent {i}: first token {r[0]:.2f}s, total {r[1]:.2f}s, {r[2]} chunks")
    print(f"TTFT median {statistics.median(ttfts):.2f}s, worst {max(ttfts):.2f}s")


if __name__ == "__main__":
    main()

uv run python agents_ttft.py --url http://mac1:8080/v1/chat/completions \
  --agents 4 --prompt-words 2000 --max-tokens 64

We could not run this script end to end. On our Linux CPU test box, the mlx-lm server stopped answering streaming requests after the first few, so we have no numbers to show. Run it on your cluster with 1, 2, 4, and 8 agents, and with one very long prompt mixed in with short ones. Watch the worst first-token time, not only the median.

Rules for a shared cluster

• Cap prompt length per request, so one agent cannot start a giant prefill for everyone (see Keep a Mac MLX box stable).
• Keep system prompts and tool lists stable so the prompt cache helps. --prompt-cache-size and --prompt-cache-bytes control how much it keeps.
• Give agents a timeout that fits the long-prompt numbers above, not a 10-second web default.
• Check tool calls and JSON from the model before you act on them (Reliable Local Agents).

Sources

• Apple Machine Learning Research: Distributed LLM inference with MLX and RDMA
• mlx-lm SERVER.md
• mlx-lm repository (server.py, benchmark.py)