feat(placement): bond parallel RDMA links in JACCL device matrix

[Nottlespike]

Motivation

get_mlx_jaccl_devices_matrix selected a single RDMA interface per device pair and discarded any additional parallel links via an early break:

for conn in cycle_digraph.get_all_connections_between(node_i, node_j):
    if isinstance(conn, RDMAConnection):
        matrix[i][j] = conn.source_rdma_iface
        break  # <-- discards extra parallel links

This left extra Thunderbolt links unused even when two Macs are connected by more than one cable. The underlying JACCL ring backend is already built to bond: RING_MAX_CONNS = 4, RingGroup accepts left_devices/right_devices as vectors, and RingImpl stripes across n_conns_ wires with per-wire buffers. The backend's JSON parser (parse_devices_json) also already accepts either a single string or an array of strings per cell. exo was the bottleneck — it threw away links the backend was designed to consume.

Changes

• src/exo/master/placement_utils.py — get_mlx_jaccl_devices_matrix now:
  • Collects every RDMAConnection interface per pair instead of breaking on the first.
  • Emits matrix cells as list[list[list[str] | None]] (a list of interface names per cell, or None on the diagonal).
  • Normalises link counts down to the minimum observed across all pairs, so the backend's ring validator (which requires every ring edge to carry the same number of connections — is_valid_ring at jaccl.cpp:130-141) never rejects an asymmetric matrix.
• src/exo/shared/types/worker/instances.py — MlxJacclInstance.jaccl_devices widened from list[list[str | None]] to list[list[list[str] | None]].
• src/exo/master/tests/test_placement.py — updated the existing tensor-parallel matrix assertions to expect list cells, plus two new tests:
  • test_jaccl_devices_matrix_bonds_parallel_rdma_links — multiple parallel links survive into the cell.
  • test_jaccl_devices_matrix_normalises_asymmetric_link_counts — asymmetric link counts are trimmed to a uniform minimum.

Why It Works

The JACCL backend (the forked mlx-jaccl-fix-small-recv wheels exo pins) consumes MLX_IBV_DEVICES as a JSON matrix. Its parser handles both forms per cell:

// jaccl.cpp:44-49
if (names.is_string()) {
  result[rank][dst].push_back(names);
} else if (names.is_array()) {
  for (auto name_it = names.begin(); name_it != names.end(); name_it++) {
    result[rank][dst].push_back(*name_it);
  }
}

RingGroup then receives vector<string> left_devices/right_devices, calls create_connections (plural), and RingImpl actively stripes sends/recvs across n_conns_ wires. So emitting arrays was always supported downstream — exo just wasn't feeding the data through.

The uniform-count normalisation is required because is_valid_ring rejects a matrix where ring edges have differing connection counts (jaccl.cpp:136). Trimming to the minimum is the conservative choice — bonding only helps when it's consistent across the ring, and a partially-bonded ring would be rejected anyway.

Known limitation: the mesh path (get_mesh_connectivity, jaccl.cpp:152) takes only devices_[rank][dst][0], so bonding is effective for ring/tensor-parallel all-reduce (where the heavy traffic is) but not for the mesh fallback. That's a JACCL limitation, not addressable from exo.

Test Plan

Manual Testing

Hardware available: Two Thunderbolt 5-capable Macs, connected by two Thunderbolt 4 cables (TB4 cables between TB5 ports — links negotiate to TB4 signalling). This is exactly the multi-link scenario this change targets: two parallel cables between one pair, which the matrix builder now feeds to the backend as a two-element cell instead of discarding one.

End-to-end throughput benchmarking on this 2×TB4-cable setup is the planned validation. The change is already verified at the configuration/data level (both interface names reach the backend in the cell) and by reading the JACCL ring striping code. A bandwidth measurement on the dual-cable pair vs single-cable would confirm the gain and is the concrete next step I can run with the hardware above.

Automated Testing

• uv run ruff check — clean
• uv run basedpyright — 0 errors, 0 warnings
• uv run pytest src/exo/master/tests/test_placement.py src/exo/shared/tests/ src/exo/utils/info_gatherer/ — 79 passed, 1 skipped (macOS-only test_tb_parsing)
• New tests cover: multi-link bonding (2 links → both in cell), asymmetric normalisation (2-link pair trimmed to match 1-link pairs), and the existing single-link case still produces correct single-element list cells.

Formatting note

CONTRIBUTING.md specifies nix fmt. I don't have Nix available in my environment, so I ran ruff format instead. If CI's treefmt config wants something different, a maintainer can re-run nix fmt on this branch.

---

Draft — opening for review, not ready to merge. Happy to adjust naming, test coverage, or the normalisation strategy based on feedback.