Cluster collectives, quantized numerical fixes and optional Metal allocation cap

docs/src/usage/unified_memory.rst

@@ -76,3 +76,21 @@ If we time the computation fully on the GPU, we get 2.8 milliseconds. But if we
 run the computation with ``d1=mx.gpu`` and ``d2=mx.cpu``, then the time is only
 about 1.4 milliseconds, about twice as fast. These times were measured on an M1
 Max.
+
+Native allocation footprint bounds
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The C++ ``get_allocation_footprint_policy()`` returns a detached value that
+bounds one allocator buffer, including alignment and permitted cache reuse.
+On Metal, every positive request has a minimum footprint of one VM page
+before the cache-reuse allowance is added. This covers the page-backed
+fallback used by the optional process-allocation cap, including requests
+smaller than one page. A policy captured before cap installation remains
+valid afterward. Zero bytes still has a zero bound; arithmetic overflow is
+refused.
+
+For example, with 16,384-byte VM pages, requests from 1 through 16,384 bytes
+have an upper bound of 32,767 bytes. The bound is conservative; it is not a
+measurement of the buffer's actual backing. The cap accounts for its fixed
+heap separately. Per-buffer footprint bounds and the optional cap do not
+replace whole-process memory budgets or OS headroom checks.

fork.yaml

@@ -106,16 +106,33 @@ def:
         overflow-checked) cache-reuse limit shared with `BufferCache`. Metal, CPU and CUDA
         allocators implement the same interface; the allocator's own behaviour stays the
         authority. These describe allocator accounting, not process RSS. (PR #15)
+        The public process-allocation cap API adds install-once Metal backing enforcement.
+        Reservations include pending allocations and the fixed heap. Owned page-aligned VM
+        regions use Metal no-copy buffers; credit returns after the deallocator unmaps them.
+        Foreign wrappers retain caller ownership. CPU/CUDA reject this optional API; soft
+        memory limits are unchanged. Residency insertion rolls back failed registration before
+        callers release buffers, including ordinary uncapped allocations. Standalone CPU and
+        real-Metal tests cover delayed release, churn, and injected construction failures.
+        Metal footprint quotations apply a one-VM-page minimum to every positive request
+        before adding the cache-reuse allowance, covering capped page-backed fallback even
+        when the policy was captured before cap installation. Standalone fresh-process
+        ordinary and capped cases compare actual backing with detached and direct bounds
+        at subpage and page boundaries; zero and overflow remain explicit controls.
+        The native footprint contract is documented in `docs/src/usage/unified_memory.rst`.
       globs:
         - "mlx/memory.h"
         - "mlx/backend/common/allocation_footprint.h"
+        - "mlx/backend/common/process_allocation_cap.h"
+        - "tests/standalone/process_allocation_cap*.cpp"
         - "mlx/backend/common/buffer_cache.h"
         - "mlx/backend/metal/allocator.h"
         - "mlx/backend/metal/allocator.cpp"
+        - "mlx/backend/metal/resident.cpp"
         - "mlx/backend/cuda/allocator.h"
         - "mlx/backend/cuda/allocator.cpp"
         - "mlx/backend/no_gpu/allocator.cpp"
         - "tests/allocator_tests.cpp"
+        - "docs/src/usage/unified_memory.rst"
     - title: "FP32 partial sums in two-pass vector SDPA"
       description: |
         The two-pass vector SDPA cast unnormalized partial sums to the input dtype before
@@ -234,15 +251,61 @@ def:
         - "docs/src/usage/environment_variables.rst"
         - "tests/jaccl/*"
         - "tests/jaccl/fake_verbs/infiniband/verbs.h"
+    - title: "Two-rank sealed all-sum"
+      description: |
+        The embedding program installs an authenticated frame codec for a two-rank
+        floating-point all-sum. The CPU communication stream seals each partial,
+        gathers the frames, opens the peer frame and adds in rank order. Missing
+        codecs and authentication failures refuse evaluation without a plaintext
+        fallback. The distributed primitive retains its temporary arrays until
+        the queued work completes.
+      globs:
+        - "mlx/distributed/CMakeLists.txt"
+        - "mlx/distributed/primitives.h"
+        - "mlx/distributed/sealed.cpp"
+        - "mlx/distributed/sealed.h"
+    - title: "Quantized tile bounds and distinct loader symbols"
+      description: |
+        Sorted gathered matrix products use the remaining reduction width for a
+        partial activation tile. Safe affine and floating-quantized weight loaders
+        check row and packed-column bounds on both legacy and NAX paths before
+        reading weights or scales, then
+        zero invalid lanes. This prevents finite tails from reading the next row,
+        including a following NaN row. Exact Python regressions cover affine,
+        MXFP4, MXFP8 and directly packed NVFP4 with FP16 and BF16 inputs.
+        Each affine/floating-point, legacy/NAX loader has a distinct type name.
+        Their layouts or bodies differ; sharing a name corrupts MXFP4 results
+        when the offline Metal linker combines the translation units. The
+        incompatible affine/FP pack-factor helpers and legacy/NAX vector loaders
+        also use distinct names.
+        Aligned dispatch and accumulation precision remain unchanged.
+      globs:
+        - "mlx/backend/metal/kernels/quantized.h"
+        - "mlx/backend/metal/kernels/quantized_nax.h"
+        - "mlx/backend/metal/kernels/fp_quantized.h"
+        - "mlx/backend/metal/kernels/fp_quantized_nax.h"
+        - "python/tests/test_quantized.py"
+    - title: "CPU quantized accumulation precision"
+      description: |
+        Scalar affine, MXFP4, MXFP8 and NVFP4 matrix products accumulate in FP32
+        and convert once when writing each output. Both matrix orientations and
+        the FP group sum preserve partial sums instead of rounding every addition
+        to BF16 or FP16. Exact CPU regressions compare constant products with
+        dequantized matmul, including sorted and unsorted gathered affine paths.
+        Quantization, packing, indexing and Metal kernels remain unchanged.
+      globs:
+        - "mlx/backend/cpu/quantized.cpp"
+        - "python/tests/test_quantized.py"
     - title: "Upstream v0.32.2 sync carried in PR #10"
       description: |
         **Not fork work.** PR #10 moved the fork from an 0.32.0-era upstream main to official
         v0.32.2 by applying upstream's changes rather than merging them, so `git` sees them as
-        fork commits even though every file below is byte-identical to upstream `v0.32.2`
+        fork commits even though the imported content matches upstream `v0.32.2`
         (CUDA backend, CPU SIMD, kernels, Python bindings and tests, repository CI and
         templates, `AGENTS.md` / `CLAUDE.md`, docs). Listed here so the page does not credit
         Layr-Labs with upstream's changes; the section disappears once the fork rebases onto
-        an upstream commit at or past v0.32.2.
+        an upstream commit at or past v0.32.2. Later fork changes to these paths
+        are described in their own sections, including quantized NAX tile bounds.
       globs:
         - ".pre-commit-config.yaml"
         - "AGENTS.md"

mlx/backend/common/process_allocation_cap.h

@@ -0,0 +1,192 @@
+// Copyright © 2026 Eigen Labs.
+#pragma once
+
+#include <cstddef>
+#include <limits>
+#include <memory>
+#include <mutex>
+#include <stdexcept>
+
+namespace mlx::core::allocator {
+
+struct ProcessAllocationCapSnapshot {
+  bool installed{false};
+  size_t limit_bytes{0};
+  size_t heap_bytes{0};
+  size_t live_bytes{0};
+  size_t cached_bytes{0};
+  size_t pending_bytes{0};
+  size_t unreclaimed_bytes{0};
+  size_t charged_bytes{0};
+  size_t peak_bytes{0};
+};
+
+// Reference release alone cannot refund backing credit.
+class ProcessAllocationCap {
+ public:
+  void install(size_t limit, size_t heap) {
+    std::lock_guard lock(mutex_);
+    if (state_.installed || !limit || heap > limit) {
+      throw std::invalid_argument("allocation cap install refused");
+    }
+    state_.installed = true;
+    state_.limit_bytes = limit;
+    state_.heap_bytes = heap;
+    state_.charged_bytes = state_.peak_bytes = heap;
+  }
+
+  void reserve(size_t bytes) {
+    std::lock_guard lock(mutex_);
+    if (!state_.installed ||
+        bytes > state_.limit_bytes - state_.charged_bytes) {
+      throw std::runtime_error("native allocation cap exceeded");
+    }
+    state_.pending_bytes += bytes;
+    state_.charged_bytes += bytes;
+    if (state_.charged_bytes > state_.peak_bytes) {
+      state_.peak_bytes = state_.charged_bytes;
+    }
+  }
+
+  // Only a reservation for which no backing was acquired can be rolled back.
+  void failed_before_acquisition(size_t bytes) {
+    std::lock_guard lock(mutex_);
+    require(state_.pending_bytes, bytes);
+    state_.pending_bytes -= bytes;
+    state_.charged_bytes -= bytes;
+  }
+  void acquired(size_t bytes) {
+    std::lock_guard lock(mutex_);
+    transfer(state_.pending_bytes, state_.live_bytes, bytes);
+  }
+  void cache(size_t bytes) {
+    std::lock_guard lock(mutex_);
+    transfer(state_.live_bytes, state_.cached_bytes, bytes);
+  }
+  void reuse(size_t bytes) {
+    std::lock_guard lock(mutex_);
+    transfer(state_.cached_bytes, state_.live_bytes, bytes);
+  }
+  void relinquish(size_t bytes, bool cached = false) {
+    std::lock_guard lock(mutex_);
+    transfer(
+        cached ? state_.cached_bytes : state_.live_bytes,
+        state_.unreclaimed_bytes,
+        bytes);
+  }
+  void destroyed(size_t bytes) {
+    std::lock_guard lock(mutex_);
+    require(state_.unreclaimed_bytes, bytes);
+    state_.unreclaimed_bytes -= bytes;
+    state_.charged_bytes -= bytes;
+  }
+  ProcessAllocationCapSnapshot snapshot() const {
+    std::lock_guard lock(mutex_);
+    return state_;
+  }
+
+ private:
+  static void require(size_t available, size_t bytes) {
+    if (bytes > available) {
+      throw std::logic_error("allocation cap ledger mismatch");
+    }
+  }
+  static void transfer(size_t& from, size_t& to, size_t bytes) {
+    require(from, bytes);
+    if (bytes > std::numeric_limits<size_t>::max() - to) {
+      throw std::overflow_error("allocation cap ledger overflow");
+    }
+    from -= bytes;
+    to += bytes;
+  }
+  mutable std::mutex mutex_;
+  ProcessAllocationCapSnapshot state_;
+};
+
+// Shared by allocator bookkeeping and the Metal deallocator block. The callback
+// never takes the allocator mutex. Release must finish before credit is
+// returned.
+class ProcessBackingLease {
+ public:
+  using Release = bool (*)(void*, size_t);
+  ProcessBackingLease(std::shared_ptr<ProcessAllocationCap> cap, size_t bytes)
+      : cap_(std::move(cap)), bytes_(bytes) {
+    cap_->reserve(bytes_);
+  }
+  ~ProcessBackingLease() {
+    finish();
+  }
+  void attach(void* pointer, Release release) noexcept {
+    std::lock_guard lock(mutex_);
+    pointer_ = pointer;
+    release_ = release;
+  }
+  void acquired() {
+    std::lock_guard lock(mutex_);
+    if (state_ != State::pending) {
+      throw std::logic_error("backing already retired");
+    }
+    cap_->acquired(bytes_);
+    state_ = State::live;
+  }
+  void cache() {
+    std::lock_guard lock(mutex_);
+    if (state_ != State::live) {
+      throw std::logic_error("backing cache mismatch");
+    }
+    cap_->cache(bytes_);
+    state_ = State::cached;
+  }
+  void reuse() {
+    std::lock_guard lock(mutex_);
+    if (state_ != State::cached) {
+      throw std::logic_error("backing reuse mismatch");
+    }
+    cap_->reuse(bytes_);
+    state_ = State::live;
+  }
+  void retire() {
+    std::lock_guard lock(mutex_);
+    if (state_ != State::live && state_ != State::cached) {
+      throw std::logic_error("backing retirement mismatch");
+    }
+    cap_->relinquish(bytes_, state_ == State::cached);
+    state_ = State::retired;
+  }
+  void finish() noexcept {
+    std::lock_guard lock(mutex_);
+    if (state_ == State::finished) {
+      return;
+    }
+    // Foreign wrappers have no release function: their caller retains
+    // ownership.
+    const bool released = !release_ || release_(pointer_, bytes_);
+    if (state_ == State::pending && released) {
+      cap_->failed_before_acquisition(bytes_);
+    } else {
+      if (state_ == State::pending) {
+        cap_->acquired(bytes_);
+      }
+      if (state_ != State::retired) {
+        cap_->relinquish(bytes_, state_ == State::cached);
+      }
+      if (released) {
+        cap_->destroyed(bytes_);
+      }
+    }
+    // A failed backing release is permanently charged; it cannot be retried
+    // after ownership has been relinquished to a deallocator callback.
+    state_ = State::finished;
+  }
+
+ private:
+  enum class State { pending, live, cached, retired, finished };
+  std::shared_ptr<ProcessAllocationCap> cap_;
+  size_t bytes_;
+  void* pointer_{nullptr};
+  Release release_{nullptr};
+  State state_{State::pending};
+  std::mutex mutex_;
+};
+
+} // namespace mlx::core::allocator

mlx/backend/cpu/quantized.cpp

@@ -108,15 +108,16 @@ void _qmm(
   constexpr int bytes_per_pack = get_bytes_per_pack(bits);
   constexpr int packs_in_group = group_size / pack_factor;
 
+  std::vector<float> sums(N);
   for (int m = 0; m < M; m++) {
     const uint8_t* w_local = (const uint8_t*)w;
     const T* scales_local = scales;
     const T* biases_local = biases;
 
-    std::fill(result, result + N, 0);
+    std::fill(sums.begin(), sums.end(), 0);
 
     for (int k = 0; k < K; k++) {
-      T* result_local = result;
+      float* result_local = sums.data();
       T xi = *x++;
 
       for (int n = 0; n < N; n += group_size) {
@@ -147,6 +148,7 @@ void _qmm(
       }
     }
 
+    std::copy(sums.begin(), sums.end(), result);
     result += N;
   }
 }
@@ -174,7 +176,7 @@ void _qmm_t(
 
     for (int n = 0; n < N; n++) {
       const T* x_local = x;
-      T sum = 0;
+      float sum = 0;
       for (int k = 0; k < K; k += group_size) {
         T scale = *scales_local++;
         T bias = *biases_local++;
@@ -458,14 +460,15 @@ void fp_qmm(
   constexpr int pack_factor = get_pack_factor(bits, 8);
   constexpr int packs_in_group = group_size / pack_factor;
 
+  std::vector<float> sums(N);
   for (int m = 0; m < M; m++) {
     const uint8_t* w_local = (const uint8_t*)w;
     const uint8_t* scales_local = scales;
 
-    std::fill(result, result + N, 0);
+    std::fill(sums.begin(), sums.end(), 0);
 
     for (int k = 0; k < K; k++) {
-      T* result_local = result;
+      float* result_local = sums.data();
       T xi = *x++;
 
       for (int n = 0; n < N; n += group_size) {
@@ -484,6 +487,7 @@ void fp_qmm(
         }
       }
     }
+    std::copy(sums.begin(), sums.end(), result);
     result += N;
   }
 }
@@ -506,11 +510,11 @@ void fp_qmm_t(
 
     for (int n = 0; n < N; n++) {
       const T* x_local = x;
-      T sum = 0;
+      float sum = 0;
       for (int k = 0; k < K; k += group_size) {
         T scale = dequantize_scale<T, group_size>(*scales_local++);
 
-        T gsum = 0;
+        float gsum = 0;
         for (int kw = 0; kw < packs_in_group; kw++) {
           if constexpr (bits == 4) {
             gsum += (*x_local++) * static_cast<T>(FP4_LUT[w_local[0] & 0xf]);