[AlexCheema]

Stacked on #2507 (the per-family folder refactor); the move to upstream mlx-lm goes on top of this.

What

platform	before	after
macOS	MLX built from git 0e3ff36 (2026-10-03, after 0.32.3)	MLX built from the v0.32.4 tag (36969c0, 2026-10-14), 56 commits later
Linux	mlx 0.32.0 and mlx-cuda-12/13 0.32.0 wheels from a GitHub release, built from upstream main in May; mlx-cpu 0.31.2 from PyPI	mlx, mlx-cpu, mlx-cuda-12 and mlx-cuda-13 0.32.4 from PyPI

Why

• Linux needs MLX >= 0.32.3 for upstream mlx-lm. Its quantized MoE layers pass global_scale to
  mx.gather_qmm, and the 0.32.0 wheels don't take that argument, so every quantized MoE would fail. The release
  wheels were a stopgap until PyPI had CUDA builds. The JACCL small-receive fix they carried has been upstream
  since Handle data smaller than BUFFER_SIZE in jaccl recv ml-explore/mlx#3033.
• macOS: v0.32.4 adds the fix for 32-bit dispatch grid overflow in Metal reductions (#4437). Sparse-attention
  indexers hit it at long context (GLM-5.x recall). It also adds the Metal fence deadlock fix (#4689) and keeps
  split-K quantized matmul partials in float32 (#4641).
• Building from the tag (rather than the PyPI macOS wheel) keeps the nix and PyInstaller builds as they are. The nix
  patch for MLX's CMake files is regenerated: v0.32.4 compiles extension Metal sources to AIR before linking
  (cmake/extension.cmake), and the old hunk no longer applied.

Verification

On a Mac Studio (M3 Ultra, 96 GB), this branch was compared with its base (MLX 0e3ff36, same exo code). Teacher-forced float32 log-probs came from exo's
own loading and sharding (single device, 2x tensor parallel, 2x pipeline parallel; prefill plus 16 decode steps),
against Hugging Face transformers (PyTorch, CPU) on the same weights:

model	single	2x TP	2x PP	KL to reference
Llama-3.2-1B-Instruct (bf16)	identical	identical	identical	2.8e-4
Qwen3.5-0.8B (bf16, Gated DeltaNet)	identical	identical	identical	4.2e-4
gpt-oss-20b (MXFP4 experts, bf16 attention)	identical	identical	identical	1.0e-3
gpt-oss-20b-MXFP4-Q8 (8-bit affine attention)	changed	changed	changed	see below

The 8-bit affine model changes: 0.32.4 reworked quantized matmul (float32 split-K partials, new thin-matrix
kernels). Against a float32 run of the same weights (affine layers dequantized), KL is 5.8e-3 -> 7.7e-3 single
device (paired difference +1.7e-3, 95% interval -2e-4 to +4.8e-3) and 6.3e-3 -> 5.7e-3 TP. There is no significant
shift either way, and MXFP4 experts are untouched.

• Linux CPU (an aarch64 Linux machine; no CUDA machine was free):
  • uv sync --locked --extra mlx-cpu installs mlx and mlx-cpu 0.32.4, and mx.gather_qmm accepts
    global_scale.
  • A 4-bit quantized Qwen3-MoE runs and decodes.
  • Llama-3.2-1B-Instruct-4bit through exo's batch generator completes "The capital of France is" with
    "Paris. The capital of Germany is Berlin. …".
  • The same quantized MoE also runs with upstream mlx-lm on top (the next PR). Under MLX 0.32.0 it fails:
    TypeError: gather_qmm(): incompatible function arguments.
  • The Linux CPU extra on main doesn't import at all: it pairs the 0.32.0 mlx wheel with mlx-cpu 0.31.2
    (undefined symbol … linalg3det). This fixes that too.
• macOS: uv sync --extra mlx builds MLX 0.32.4 from source; nix build .#exo-test-env builds and runs on the GPU;
  mlx/lib (libmlx, mlx.metallib) is in place for PyInstaller. The Linux nix packages were not built here.
• basedpyright, ruff check, nix fmt, uv lock --check and pytest (1842 passed, 3 skipped) pass.

Lockfile note

Relocking with current uv also resolves torch for Linux. The existing lock had torch only for macOS, although
pyproject.toml asks for it on Linux too. Torch now comes from the CUDA 13 index for every Linux extra, including
mlx-cpu. Relocking main with nothing else changed gives the same torch delta, so this is not from MLX. The torch
index markers deserve their own look.

🤖 Generated with Claude Code