super-kaiju / SERVING.md
restokes92's picture
Upload SERVING.md with huggingface_hub
160e22d verified
|
Raw
History Blame Contribute Delete
2.16 kB

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 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.