feat: single-node multi-GPU tensor parallelism to load and serve models larger than one GPU (CUDA/DGX)

[inureyes]

Goal

Make it possible to load and serve a model whose weights exceed a single GPU's memory by sharding those weights across multiple physical GPUs on one node (for example a DGX-class box with 8 GPUs over NVLink), with each tensor-parallel rank's weights and compute resident on its own GPU and real cross-GPU collectives between ranks.

The headline deliverable: mlxcel-server --tp-size N (and mlxcel generate --tp-size N) loads a model that does NOT fit on a single GPU and serves it correctly across N GPUs on one node.

Current state

Tensor parallelism is extensively implemented (shard planning, ring collectives, parallel attention/FFN/MoE shape math, sharded loading, TP scheduler/executor/barrier, per-rank KV cache) and an in-process TP runtime is wired into the load path:

• src/loading/mod.rs:431 load_model_with_tensor_parallel constructs TensorParallel*Model when tp_size > 1 (supported: Llama/Qwen2, Qwen3, Qwen3.5, Gemma3, Gemma4, Ernie4.5, HunyuanV1Dense).
• Reachable from CLI generate (src/commands/generate.rs:130) and the server (src/server/model_worker.rs:234,753).

However the runtime is a single-device numerical emulation, not real multi-GPU execution:

• src/distributed/tensor_parallel/llama_runtime.rs:98-148: the forward iterates ranks with a sequential .iter().map() on the default device and sums partials locally via reduce_sum (llama_runtime.rs:2288). No threads, no per-rank device, no network/NVLink collective.
• src/distributed/tensor_parallel/llama_runtime.rs:76-95 from_full_weights: every rank's sharded weights are materialized on the single default device, so aggregate memory is not reduced and a model larger than one GPU still cannot be loaded.

Root cause in the bindings: the device API is binary CPU/GPU only, with no device index.

• src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:3780 set_default_device(bool gpu) maps to Device::gpu (index 0).
• src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:18 new_stream_on_device(bool gpu) likewise.
• src/lib/mlxcel-core/src/lib.rs:45,54,1708 expose only the boolean forms.

So even on a CUDA box with many GPUs, all work runs on the one default GPU.

MLX itself is fetched from upstream ml-explore/mlx at GIT_TAG a6ec7123dac814417147e21d4aeed694924ddd4d with a CUDA backend and custom CUDA patches (src/lib/mlx-cpp/CMakeLists.txt:92-95). mlx::core::Device carries a DeviceType plus an integer index, and MLX provides a distributed module (MPI/ring, and NCCL on CUDA), but it is process-based and not currently bound in mlxcel-core.

Decomposition

This epic is delivered by two dependent sub-issues:

1. Sub-issue: device-index-aware device and stream bindings in mlxcel-core (the foundation that lets us target GPU 1..N and query device count).
2. Sub-issue: multi-GPU tensor-parallel runtime that places each rank's weights and compute on its own GPU and reduces partials with real cross-GPU collectives. Depends on (1).

Sub-issues

Phase 1

• feat: device-index-aware device and stream bindings in mlxcel-core (multi-GPU foundation) #487

Phase 2

• feat: multi-GPU tensor-parallel runtime with per-rank device placement and cross-GPU collectives #488 (depends on feat: device-index-aware device and stream bindings in mlxcel-core (multi-GPU foundation) #487)

Cross-cutting feasibility spike (gating)

Before committing to the intra-process design in sub-issue 2, verify that MLX's CUDA backend at the pinned commit can, within a single process: (a) allocate arrays on a non-zero GPU index, (b) run ops on a non-default device, and (c) move arrays between devices (peer copy). If intra-process multi-GPU is not supported by MLX as pinned, fall back to a multi-process design (one process per GPU on localhost/NVLink, reusing the existing distributed transport path), or evaluate bumping the MLX pin. The spike outcome decides sub-issue 2's architecture and must be recorded in that issue.

Integration and validation (epic acceptance)

The epic is complete only when the capability is integrated end to end, not just present as modules:

• gpu_device_count() (from sub-issue 1) is surfaced at server/CLI startup (a diagnostic log line such as "detected N GPUs").
• load_model_with_tensor_parallel places shard r on GPU r; a memory probe shows each GPU holds approximately 1/N of the weights.
• mlxcel generate --tp-size N and mlxcel-server --tp-size N run on N GPUs and produce logits matching the single-GPU replicated reference within tolerance.
• A model that does not fit on one GPU loads and serves with tp-size N (the DGX goal), validated on real multi-GPU hardware.
• Clear, actionable errors when tp_size > device_count(), and a documented behavior on Metal/Apple (single GPU): tp_size > 1 keeps the current in-process emulation or errors clearly.
• Docs updated (multi-GPU usage, flags, hardware requirements, limitations).

Non-goals

• Multi-node scale-out already exists via pipeline parallelism over TCP/Thunderbolt/RDMA and disaggregated prefill/decode; this epic is about a single node.
• Metal/Apple Silicon multi-GPU is not applicable (one GPU per machine, unified memory).
• 2D pipeline-plus-TP composition is tracked separately.

References

• Bindings device API: src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:18,3780; src/lib/mlxcel-core/src/lib.rs:45,54,1708
• TP runtime: src/distributed/tensor_parallel/llama_runtime.rs:76-148,2288
• TP load wiring: src/loading/mod.rs:431-476
• MLX pin: src/lib/mlx-cpp/CMakeLists.txt:92-95
• CUDA feature: Cargo.toml:48

[inureyes]

Progress: sub-issue #487 complete; #488 deferred to CUDA hardware

Sub-issue 1 of 2 is done. #487 (device-index-aware device and stream bindings) merged via PR #505 (38dbaad), implemented and verified end to end on macOS / Apple Silicon.

What landed (the foundation):

• Portable bridge surface in mlxcel-core: gpu_device_count(), new_stream_on_gpu_index, set_default_device_index, new_thread_local_stream_gpu_index, and copy_array_to_stream.
• Safe Rust wrappers (new_stream_on_gpu, set_default_gpu_device, new_thread_local_stream_on_gpu) that validate the index against the device count and return a typed InvalidGpuIndex error.
• Startup integration: the server banner and the CLI model-load path both log the detected GPU count, so the new API is exercised in a real path.

Design note: the foundation needs no CUDA-specific code. MLX at the pinned commit already exposes a portable device_count(DeviceType) and the Device(DeviceType, index) constructor, so the same code compiles on Metal (1 GPU), a CUDA multi-GPU host, and a CPU-only build with no CUDA headers and no preprocessor gating. This is a simplification over the cudaGetDeviceCount path the issue anticipated.

#488 (multi-GPU tensor-parallel runtime) is deferred. Its gating feasibility spike (intra-process multi-GPU on MLX's pinned CUDA backend) and every acceptance test (logits matching across N GPUs, a model too large for one GPU loading and serving) intrinsically require a CUDA/DGX box. The current development machine is macOS / Apple Silicon (single unified-memory GPU, no CUDA toolkit), so #488 cannot be built, run, or validated here. It will be implemented and validated on CUDA hardware, consuming the bindings #487 added.

This epic stays open until #488 lands and the epic-level acceptance criteria (per-GPU shard placement, --tp-size N on N GPUs, a too-large model loading) are validated on real multi-GPU hardware.