| # Serving Super Kaiju |
|
|
| Qwen3-Coder-480B-A35B at 4-bit is ~270 GB, too big for a single 256 GB machine. Super Kaiju runs it pipeline-parallel across two Apple M3 Ultra Mac Studios (256 GB each) over Thunderbolt with RDMA, using MLX. |
|
|
| ## The three things that made it work |
|
|
| 1. **Pipeline parallelism, not tensor parallelism.** The model's layers are split across the two nodes (each node holds and runs only its slice), using mlx-lm's `PipelineMixin`. Tensor-parallel sharding is impractical for this on Apple Silicon today. |
|
|
| 2. **An uneven layer split.** A naive 50/50 split puts the whole generation working-set transient onto the node that holds the embedding and first layers, and it runs out of memory. Giving that node *fewer* layers (here a 42/20 split of the 62 layers) balances the peak: roughly 183 GB on one node, 88 GB on the other, both well under the ceiling. |
|
|
| 3. **A warmup pass.** The first cold forward compiles a large Metal command buffer that can blow past the memory ceiling. Compiling the shaders incrementally (a few layers at a time) before the first real request keeps the peak bounded. |
|
|
| ## Why a custom server |
|
|
| The off-the-shelf `mlx_lm.server` runs generation in a worker thread, but MLX binds the distributed collectives to the main thread's GPU stream, so the first real request aborts with `no Stream(gpu, 0) in current thread`. |
|
|
| [`serve/kaiju_serve.py`](serve/kaiju_serve.py) fixes this by keeping everything GPU-touching on the main thread: |
|
|
| - The lead node runs a single-threaded HTTP server, so the request handler executes on the main thread that owns the GPU stream. |
| - On each request it tokenizes the prompt, broadcasts the token ids to the other node via an `all_sum` collective, and then both nodes run the identical `stream_generate` in lockstep. The pipeline send/recv/all_gather collectives keep them in step, and a synced sampling seed keeps their sampling identical. |
| - The second node runs a main-thread worker loop that blocks in the broadcast collective between requests and participates when the lead node initiates one. |
| |
| Measured: correct code generation at ~13 tok/s, memory stable, stable across repeated requests. |
| |
