Single-Mac GLM-5.3-Flash native fixes without JACCL

docs/src/dev/custom_metal_kernels.rst

@@ -466,3 +466,17 @@ We can then implement the backwards pass as follows:
 There's an even larger speed up for the vjp:
 
 ``676.4ms -> 16.7ms => 40x speed up``
+
+Threadgroup compiler capacity
+-----------------------------
+
+Custom Metal kernels require positive threadgroup axes, nonnegative grid axes,
+a thread count within UINT32_MAX, and axes within the device limits. Invalid
+geometry is refused before output allocation. A capable original pipeline is
+kept. For a nonempty grid containing complete 512- or 1024-wide groups, exactly
+``(512, 1, 1)`` or ``(1024, 1, 1)`` may retry an insufficient pipeline with the MSL
+``max_total_threads_per_threadgroup(requestedWidth)`` declaration attribute. The retry
+uses a distinct source/library identity and must itself support the full
+threadgroup. Compile errors and remaining capacity failures propagate. This
+does not change shader math, precision, group/grid axes, or guarantee that a
+particular shader can compile at the requested width. Other geometries receive no retry.

docs/src/usage/unified_memory.rst

@@ -76,3 +76,17 @@ 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.
+Zero bytes has a zero bound; arithmetic overflow is refused. The bound does
+not apply a one-page minimum to requests smaller than one page.
+
+Capped no-copy backing is touched once per VM page from the CPU before it is
+wrapped as a Metal buffer. This makes the process physical footprint include
+pages first written by GPU work. It preserves the allocator cap and does not
+replace load admission or operating-system memory headroom. The fresh-process
+``gpu-footprint`` regression checks GPU-written backing; it requires Metal.

fork.yaml

@@ -113,16 +113,30 @@ 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.
+        Capped no-copy pages are first touched on the CPU so process physical
+        footprint accounts for pages later written by the GPU.
+        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
@@ -184,15 +198,60 @@ def:
         - "mlx/backend/metal/kernels/binary_ops.h"
         - "tests/ops_tests.cpp"
         - "python/tests/test_ops.py"
+    - title: "Validated custom Metal threadgroups and bounded compiler retry"
+      description: |
+        Custom kernels refuse invalid axes, counts beyond UINT32_MAX and axes beyond
+        the device limits before allocating outputs. A nonempty grid with complete
+        x-axis multiples of the group width and groups (512,1,1) or (1024,1,1) may retry an insufficient
+        pipeline with max_total_threads_per_threadgroup(requestedWidth) on its unique generated
+        declaration. The original capable pipeline is retained; retry compilation
+        and the selected pipeline's actual capacity remain authoritative. Other
+        geometries refuse without clamping. Shader math and dispatch axes do not change.
+      globs:
+        - "python/tests/test_metal_threadgroup_bounds.py"
+        - "docs/src/dev/custom_metal_kernels.rst"
+    - 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"
@@ -311,6 +370,15 @@ def:
         floating-point kernels have no global-scale input, and the NAX kernels keep the
         fork's bound for the last K tile. The section disappears once the fork rebases onto
         an upstream commit at or past `90089440`.
+
+        One upstream change in these NAX kernels is reverted on purpose: the simdgroup skip
+        of [ml-explore/mlx#4352](https://github.com/ml-explore/mlx/pull/4352) (`c7ff35d9`,
+        which arrived with the v0.32.2 changes; `#4572` kept it as
+        `sg_active = sgp_sm > 0`). `fp_gather_qmm_rhs_nax` and `affine_gather_qmm_rhs_nax` run
+        the MMA in every simdgroup again. With the skip, mixture-of-experts serving on a
+        two-Mac JACCL pair with one NAX GPU lost all-reduce completions and left 92 to 96 GiB
+        wired until a reboot. Stores stay bounded by `sgp_sm`, so the output is unchanged.
+        Keep the revert at every rebase onto upstream.
       globs:
         - ".gitignore"
         - "benchmarks/python/swiglu_bench.py"

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