feat: multi-GPU tensor-parallel runtime with per-rank device placement and cross-GPU collectives

[inureyes]

Part of #486.
Depends on #487.

Problem

The in-process TP runtime shards weights correctly but materializes every rank on the single default GPU and reduces partial outputs locally, so there is no memory scaling and no real parallelism. A model larger than one GPU still cannot be loaded.

• src/distributed/tensor_parallel/llama_runtime.rs:76-95 from_full_weights: all ranks built on the default device.
• src/distributed/tensor_parallel/llama_runtime.rs:98-148: forward runs ranks sequentially with .iter().map() and sums via reduce_sum (llama_runtime.rs:2288), all on one device.

Goal: place rank r's sharded weights and compute on GPU r, and replace the local sum with a real cross-GPU collective, so aggregate weights can exceed a single GPU and the forward runs across GPUs.

Feasibility spike (do this first; it gates the design)

Confirm, within a single process on the pinned MLX (ml-explore/mlx @ a6ec7123, src/lib/mlx-cpp/CMakeLists.txt:92-95): (a) arrays can be allocated on a non-zero GPU index, (b) ops can run on a non-default device, (c) arrays can be moved between devices (peer copy). Record the outcome in this issue.

• If intra-process multi-GPU works: implement the intra-process design below.
• If it does not: fall back to a multi-process design (one worker process per GPU on localhost/NVLink, reusing the existing distributed transport and collective scaffolding in src/distributed/), or propose bumping the MLX pin. Document the decision before writing the runtime.

Design (intra-process path)

Weight placement (memory scaling, the core win):

• Build rank r's model with its sharded weights resident on GPU r. Wrap the per-rank construction in from_full_weights so weight materialization runs on a device-bound (thread-local) stream for index r, using the sub-issue 1 API. Verify each GPU holds approximately 1/N of the weights.

Forward execution:

• Run each rank's per-layer attention and FFN on a worker pinned to GPU r (per-rank thread, each owning a device-bound thread-local stream; reuse the existing MlxThreadLocalStream pattern already used by the batch scheduler).
• Keep the residual stream h replicated on every device (standard Megatron-style TP). Broadcast h to all devices before the attention and FFN regions.
• Replace reduce_sum of partial attention (after o_proj) and partial FFN (after down_proj) with a cross-GPU all-reduce. The existing shape metadata already marks where reduction is required (requires_allreduce_after_o_proj, CommPattern).

Collective backend abstraction:

• Define a CrossDeviceCollective trait with two implementations:
  1. Portable copy-and-reduce fallback: gather partials onto a root device using the sub-issue 1 copy-to-device primitive, sum, then broadcast back. Correct first, slower.
  2. NCCL-backed all-reduce on CUDA for performance, by binding MLX's distributed/NCCL path or NCCL directly. Ship the fallback first; add NCCL as an optimization.

Reuse existing scaffolding where possible: TPBarrier, synchronized.rs step decisions, split_rank_caches (already per-rank), and the per-rank TPCacheManager. KV caches must be device-resident per rank.

Acceptance criteria

• Spike outcome recorded and architecture chosen (intra-process or multi-process).
• load_model_with_tensor_parallel (src/loading/mod.rs:431) places shard r on GPU r; a memory probe confirms roughly 1/N weights per GPU.
• mlxcel generate --tp-size N and mlxcel-server --tp-size N execute across N GPUs and produce logits matching the single-GPU replicated reference within tolerance (extend the existing TP parity tests, for example tests/tp_e2e.rs, tests/tensor_parallel_real_models.rs).
• A model that does not fit on a single GPU loads and serves with tp-size N on real multi-GPU hardware (the DGX goal).
• tp_size > gpu_device_count() produces a clear, actionable error.
• Documented behavior on Metal/Apple (one GPU): tp_size > 1 keeps the current in-process emulation or errors clearly; no silent wrong results.
• At least the currently supported TP model families (Llama/Qwen2, Qwen3, Qwen3.5, Gemma3, Gemma4, Ernie4.5, HunyuanV1Dense) work on the multi-GPU path, or the gap is explicitly documented.
• Docs updated; no clippy/fmt regressions; cross-platform checks pass.

Risks and fallback

• MLX intra-process multi-GPU maturity is the main risk; the spike gates the design, with a multi-process fallback over the existing distributed transport.
• Cross-device synchronization correctness (barriers around each all-reduce) must be tested under concurrency.

Out of scope

• Vocab-parallel embeddings and LM head for very large vocabularies (current default is replicated); track as a follow-up if needed.
• 2D pipeline-plus-TP composition (tracked separately).
• MoE expert-parallel multi-GPU placement beyond the listed dense families (follow-up).