Kimi-K3 (2.78T) in TP4 on four M3 Ultras — dispatch-bound decode, and some things I had wrong [ongoing debugging]

[guruswami-ai]

Status: ongoing debugging. Posting mid-investigation because the performance attribution and the failure modes look useful to others now, and because a couple of the findings contradict things I had believed for months. There is an unresolved hang described at the bottom — I am not claiming it is an MLX bug, and I am still eliminating causes on my own hardware.

---

Kimi-K3 (2.78 T params, 1.42 TiB at MXFP4) loads and generates in tensor parallel across four M3 Ultra Mac Studios over a Thunderbolt 5 mesh with the jaccl/RDMA backend.

Stack: mlx 0.32.0, mlx-lm 0.31.3, macOS 26.6 (25G72), 4× M3 Ultra 512 GB, AppleThunderboltRDMA 0.0.1.

Load	229.6 s, 420.8 GB resident per rank, 445.7 GB peak — identical on all four
Decode	~2.0 tok/s single stream · 21.2 tok/s aggregate at 32 concurrent
Prefill	~170 tok/s chunked
Context	full 1,048,576 tokens fits — 29 GB of KV against 42.3 GB headroom
Power	559 W mean / 977 W peak for the whole cluster · 14.7 Wh per 1000 tokens

Greedy decode produced identical token ids on every rank, and the output was coherent and correct — the two things a mis-shard breaks.

The result I think is most useful to this project

Decode at batch 1 is dispatch-bound, not bandwidth-bound.

Only 103 B of 2.78 T params are active per token; at TP4 each node reads 22.5 GB/token, which at 819 GB/s is a 27.4 ms roofline → 36 tok/s. Observed was ~500 ms/token. I measured the three obvious suspects:

suspect	measured	share
collectives	5.0 ms/token (all_sum at 26 µs for 7–14 KB over RDMA)	1 %
MoE gather (92 layers, top-16 of 896)	34.0 ms/token	7 %
fixed per-call overhead	~170 ms/token	the problem

Every module costs 0.23–0.48 ms regardless of how much work it does:

component	params	bw-ideal	measured	overhead
router gate	6.4 M	4 µs	234 µs	53×
latent down+up	51 M	35 µs	478 µs	14×
KDA attention	115 M	79 µs	426 µs	5×
MoE experts	132 M	91 µs	370 µs	4×
MLA attention	230 M	158 µs	462 µs	3×

Batching is what recovers it — 1.32 → 21.2 tok/s aggregate from 1 to 32 concurrent streams, with concurrency 4 costing 38.4 s against 36.4 s for a single request, i.e. nearly free.

Caveat: those per-module numbers were taken with mx.eval around each module, a sync the real model does not pay per-module, so ~170 ms is an upper bound and roughly 330 ms of the 500 ms remains unattributed. My guess is the ~186 collectives per token acting as pipeline-drain barriers rather than as latency, but I have not tested that and two earlier hypotheses of mine died on measurement.

Things I had wrong, in case others believe them too

• MTU is irrelevant to RDMA throughput. Identical runs at MTU 9000 and 1500: 9.05 vs 8.93 GB/s at 134 MB payloads, within noise. Obvious in hindsight — RDMA uses neither TCP/IP nor Ethernet frames. Still matters for the ring backend.
• mlx.distributed_config never touches MTU. Zero occurrences of mtu in its source, and interfaces read 9000 immediately after --auto-setup. I had "step 3 resets MTU" written down for months.
• The hostfile does not need copying to peers on 0.32.0. Verified by deleting it from every peer; jobs still ran correctly. This was a real requirement in older MLX.
• The Thunderbolt /30 link-locals carry no traffic. Control/rendezvous runs over Ethernet; the /30s are a prerequisite that stays unused. So a healthy management network is a prerequisite for an RDMA job — counterintuitive when debugging a mesh that "looks fine but won't launch".

A diagnostic that saved a lot of time

GPU utilisation reads 100 % while the GPU is doing nothing. Power is the honest signal:

total SoC	GPU util	meaning
~14 W	low	idle
~24 W	7 %	this rank has dropped out of the collective
40–50 W	100 %	stalled in a collective — spinning, not working
234–258 W	100 %	genuinely working

Why TP needs the whole model on every node

From mlx_lm/utils.py::sharded_load: pipeline parallel computes which files a rank needs and fetches only those, while tensor parallel has no file selection at all — every rank lazily loads the full path then calls model.shard(). Laziness keeps RAM at ~1/N, but every rank must be able to open every file. For this model that is 1.42 TiB × 4 = 5.7 TiB of disk for one model.

TP slices are contiguous within each tensor and safetensors supports partial reads, so "read only my rows" looks expressible. Related: #3208.

The unresolved part

A multi-chunk prefill hangs under TP: short prompts and prompts whose second chunk is a partial remainder complete, while a 9699-token prompt stops after the first full 2048-token chunk. Silent — no error, no timeout, no diagnostic.

I have excluded by re-running with each removed: concurrency, prompt-cache eviction, MLX_METAL_FAST_SYNCH (both values), PD exhaustion, mesh health (verified 3 peers on every node during the hang), Thunderbolt cables, mid-run link drops (123 samples, zero events), rank ordering, firmware/Secure Boot, RDMA transport (134 MB collectives at 9.33 GB/s), the Metal 60 s command-buffer timeout (zero hits in any log), and per-node network-service config.

The failure follows the host, not the rank index — so I am currently treating it as something on my machine and am rebuilding that node from scratch to find out. I will not file it as an issue until I can show it survives a clean install.

The part I would raise regardless: when a rank stops participating, the peers spin indefinitely with no timeout and no diagnostic. Whatever the trigger, that turns a five-minute problem into a multi-reboot one. Is a collective timeout something you would consider?

Write-up and tools

Full detail, the shard() traps (two of which produce running code with wrong output), the troubleshooting runbook, and the diagnostic tools are here:
https://github.com/guruswami-ai/mlx-benchmarks

Happy to run targeted experiments — the hardware is standing and I can iterate.