Skip unnecessary simdgroup computations for quantised MOE matmuls on NAX - #4352

* Return tuple in meshgrid (ml-explore#4229)

* Add endpoint parameter to linspace (ml-explore#4184)

Co-authored-by: Cheng <[email protected]>

* Fix vmap of partition/argpartition dropping the kth argument (ml-explore#4116)

* Fix nan_to_num replacing inf with 0 for float16 and bfloat16 (ml-explore#4222)

Co-authored-by: codeAnqiang-ma <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Fix einsum not broadcasting batch dimensions in batched tensordot (ml-explore#4125)

Co-authored-by: Cheng <[email protected]>

* Dequantize in float32 (ml-explore#4241)

* chore: Reject complex in erf and erfinv (ml-explore#4243)

* Fix cpu compilation failure of abs with uint (ml-explore#4240)

Co-authored-by: Cheng <[email protected]>

* Fix quantize matrix multiplication floor issue (ml-explore#4251)

* Only use MPI backend for world size > 1 (ml-explore#4210)

* chore: Reject complex in expm1, sigmoid and arctan2 (ml-explore#4257)

* Decompose small kernel-depth 3D convs into 2D convs (ml-explore#3785)

Co-authored-by: katlun-lgtm <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Fix Metal sort of a view with a negative stride (ml-explore#4252)

* Mirror the depth axis in the decomposed 3D conv when flipped (ml-explore#4277)

* Fix Metal row reductions on negative-stride views (ml-explore#4267)

Co-authored-by: Fu Xiaonan <[email protected]>

* [CUDA] Fix custom kernel cache collision for same name, different source (ml-explore#4273)

Co-authored-by: Cheng <[email protected]>

* Fix ops rejecting integers larger than INT32_MAX (ml-explore#4255)

Co-authored-by: Feli <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Fix var/std for complex numbers (ml-explore#4260)

* Fix int32 overflow in conv padded input and pad shapes (ml-explore#4258)

Co-authored-by: Cheng <[email protected]>

* chore: Reject complex in remainder (ml-explore#4270)

* chore: Compare the macOS SDK version as a version when gating JACCL (ml-explore#4286)

* Clamp ring socket transfers so a payload of 2 GiB or more can be sent (ml-explore#4281)

Co-authored-by: Cheng <[email protected]>

* chore: Use normalize_axis_index in split/unstack/partition/topk (ml-explore#4288)

* Remove grouped output in CI (ml-explore#4195)

* [CUDA] Fix finding cuda 13 headers in JIT compilation (ml-explore#3995)

* Refactor wheel building script (ml-explore#3818)

* Make mx.compile cache erasing thread safe (ml-explore#4248)

Co-authored-by: yentur <[email protected]>

* Add builds for free-threaded python (ml-explore#3812)

* Fix int32 overflow in concatenate/repeat/kron (ml-explore#4303)

* python: Widen list elements that do not fit in int32 to int64 (ml-explore#4305)

* Propagate CPU errors to events (ml-explore#3742)

Co-authored-by: Alessio Pollero <[email protected]>

* Fix mx.arange dtype inference overflow regression (ml-explore#4324)

* Add workflow to update pull request limit bypass list (ml-explore#4320)

* Support head dimension 72 in Metal full attention (ml-explore#4330)

* Patch bump to 0.32.2 (ml-explore#4333)

* Preserve subnormal float values when casting to bool (ml-explore#4224)

* python: Support assigning through a bare Ellipsis index (ml-explore#4314)

* Fix divmod truncating the quotient for floats (ml-explore#4108)

Co-authored-by: Cheng <[email protected]>

* Add force_fused option to scaled_dot_product_attention (ml-explore#4185)

* chore: Reject negative eps in the normalization layers (ml-explore#4312)

* Bound GGUF metadata string/array values against the file mapping (ml-explore#4212)

Co-authored-by: x14ngch3n <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Read each K/V byte once in gqa-8 decode attention (ml-explore#4077)

* Fix fft vmap and jvp for transforms over a subset of axes (ml-explore#4138)

* Fix median dropping NaN (ml-explore#4146)

* Fix the CPU scan over a size one axis with a padded stride (ml-explore#4139)

Co-authored-by: Cheng <[email protected]>

* chore: Validate the optimizer betas at construction (ml-explore#4310)

Co-authored-by: Cheng <[email protected]>

* `RMSNormVJP` backward writes a full `{n_rows, D}` `gw_temp` intermediate (ml-explore#4293)

* [Bug]: add default none value to axis parameter of the take_along_axis (ml-explore#4357)

Co-authored-by: Anastasiia Filippova <[email protected]>

* Add a fused full-attention path for head_dim 256 on NAX devices (ml-explore#3842)

Co-authored-by: Cheng <[email protected]>

* Update nanobind to 2.15.0 (ml-explore#4337)

* Skip unnecessary simdgroup computations for quantised MOE matmuls on NAX (ml-explore#4352)

* Add AI usage policy (ml-explore#4331)

Co-authored-by: Jake Bowhay <[email protected]>

* Raise cpu stream errors from synchronize (ml-explore#4338)

Co-authored-by: Cheng <[email protected]>

* chore: Validate eps in Adam at construction (ml-explore#4361)

Co-authored-by: Anastasiia Filippova <[email protected]>

* Bound winograd conv2d working set by tiling the batch (ml-explore#4102)

Co-authored-by: Cheng <[email protected]>

* Use a 32-row block in qmm_t_nax when one block covers all of M (ml-explore#4171)

* chore: Deduplicate fftshift and ifftshift (ml-explore#4318)

* Fix Log and Equal is_equivalent ignoring primitive state (ml-explore#4266)

Co-authored-by: Cheng <[email protected]>

* Stabilize reduced-precision InstanceNorm (ml-explore#4230)

* chore: Normalize negative axes in sort and argsort (ml-explore#4332)

* Clean up main thread compile cache before python interpreter shuts down (ml-explore#4373)

* chore: Check malformed jaccl hostfile that miss rdma in pairs (ml-explore#4284)

Co-authored-by: Cheng <[email protected]>

* Round mxfp8 block scales up to avoid saturation (ml-explore#4353)

Co-authored-by: Daniel Hiltgen <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Add support for the __array_namespace_info__  (ml-explore#4334)

* Stop a failed CUDA graph commit from poisoning the encoder (ml-explore#4356)

Co-authored-by: Cheng <[email protected]>

* Fix quantized kernels in JIT build (ml-explore#4372)

Co-authored-by: Cheng <[email protected]>

* Avoid zero work in stride-2 ConvTranspose3d (ml-explore#4343)

* [CUDA] Ce fused kernel (ml-explore#3947)

* Fix cpu exclusive scan for complex numbers (ml-explore#4272)

Co-authored-by: Cheng <[email protected]>

* Support Relocatable CUDA DLLs on Windows (ml-explore#4382)

* Use cast_to for fused AsType in compiled Metal kernels (ml-explore#4351)

Co-authored-by: katlun-lgtm <[email protected]>
Co-authored-by: Cheng <[email protected]>

* python: Declare DLPackCompatible protocol members as methods (ml-explore#4384)

* Fix quantizing sliced arrays (ml-explore#4381)

* Fix einsum dropping a trailing empty subscript (ml-explore#4299)

Co-authored-by: Cheng <[email protected]>

* Add script to run python tests (ml-explore#4393)

* Hold GIL in AttachedData destructor (ml-explore#4391)

* Bound Metal buffer COUNT, not just bytes, in MetalAllocator

The Metal allocator throws `[metal::malloc] Resource limit (N) exceeded`
when num_resources_ (the live+cached Metal buffer COUNT) reaches
resource_limit_ (the iogpu.rsrc_limit sysctl, default ~499000). Freed
buffers are recycled into a size-keyed cache whose only trim is by BYTES
(release_cached_buffers takes a bytes-to-free target, max_pool_size_ ~=
physical RAM). Under churn with many distinct buffer shapes (varied prompt
lengths, growing KV caches, multiple co-resident models) the cache fills
with entries never reused at that exact size, so the COUNT climbs to the
limit while byte usage stays modest and the byte trim never fires — the
process crashes mid-inference on a machine with most of its RAM free.

malloc() now also reclaims by count: when num_resources_ crosses a 90%
high-water mark of resource_limit_, it clears the (pure-reuse) buffer
cache so the count drops back to the live working set. Clearing the cache
only costs re-allocation, never correctness, so the count limit becomes
unreachable by any request mix or batching method while the existing byte
limits keep total memory bounded.

Adds get_num_resources()/get_resource_limit() to the public memory API
(metal + no_gpu + cuda backends) so the count and its ceiling are
observable from callers. Adds an MLX_RESOURCE_LIMIT env override that can
only LOWER the ceiling (clamped to the OS limit, strictly validated) to
exercise the trim deterministically and as an operator safety valve.

* perf(mlx): opt-in Gemma 4 expert-QMM tile kernel with parallel descriptor builder (#4)

* perf(mlx): add opt-in Gemma 4 expert-QMM tile kernel with parallel descriptor builder

Adds a distinctly-named expert QMM implementation for the Gemma 4
26B-A4B MoE production shapes, gated by MLX_GATHER_QMM_EXPERT_SLICES:

- qmm_t_expert_impl: BM32 expert tile body (BM16 fallback rows) taking a
  private/by-value row count; the shared qmm_t_impl constant-address ABI
  and all ordinary gathered/batched/dense QMM routes are unchanged.
- build_gemma4_sorted_expert_tiles_bm32: one 128-thread threadgroup
  replaces the reference design's single-GPU-thread serial builder;
  parallel expert-range binary search, Hillis-Steele scan, and strided
  upper-bound descriptor emission.
- Selector runs after the NAX-first route and requires affine BF16
  transposed inputs, 4-bit gs=64 weights, 128 experts, assignment counts
  of exactly 4096/8192/16384, and the exact gate/up or down rank-3
  shapes; every miss keeps the legacy route. NAX engagement is
  non-engagement, never bypassed.
- device.{h,cpp}: one-shot request resolution, nonthrowing dual-symbol
  AOT probe/prewarm, relaxed-atomic diagnostics (requested, aotAvailable,
  naxAvailable, hits, per-class fallbacks).
- gpu_tests: exact-shape arithmetic parity, fallback, and counter
  invariant probes.

Retention standing (2026-08-09 production matrix): opt-in experiment.
Standalone profile dropped (prefill -10.2% vs bracket); paired
weighted-unsort+R1 profile retained-final (prefill +1.8%, TTFT -7.5%,
decode +3.3%, arrival E2E +12.0%). NOTE: this source post-dates the
benchmarked binaries/metallib (post-measurement kernel-body edit);
rebuild and re-verify before any performance claim.

* fix(mlx): fail-safe sortedness check in gemma expert tile builder; counter/atomic hygiene

Review-wave fixes for the R1 expert-QMM path:

- N1 (sortedness trust): build_gemma4_sorted_expert_tiles_bm32 now
  verifies each thread's post-binary-search segment boundary against the
  generalized invariant indices[start - 1] < lid <= indices[start]
  (edge threads check their single neighbor), votes per simdgroup via
  simd_or, folds the votes through threadgroup memory, and on any
  violation retracts count[0] to 0 (tile kernel then early-returns) and
  records the violation in count[1]; the buffer ABI is unchanged
  (count index 1 was previously unused). try_gemma4_expert_qmm allocates
  the second count element, drains the encoder after the builder, and
  re-routes a retracted call to the order-agnostic legacy path instead of
  dispatching the tile kernel (zero count is unambiguous: the selector's
  assignment gate guarantees M is 4096/8192/16384).
- N2 (route-condition duplication): the sorted-RHS gate literal that
  appeared (negated) in the diagnostics record and in the dispatch
  decision is now the shared static constexpr predicate
  takes_sorted_rhs_route, so future tuning of the 16/4 thresholds cannot
  desynchronize counter vs route.
- N3 (per-call bias normalization): gather_qmm_rhs no longer spends
  ensure_row_contiguous on biases before classification reads the raw
  tensor's fields; normalization runs only inside the winning-route
  branch (hit semantics unchanged; the legacy block keeps its own
  normalization point and ordering).
- N4 (armed_ data race): Gemma4ExpertQMMCounters::armed_ is now
  std::atomic<bool> with relaxed loads/stores in armed(), snapshot(),
  snapshot_and_disarm() (read-then-write order preserved) and
  clear_and_arm(); the class remains non-copyable, now enforced.

* fix(mlx): make the R1 sortedness fail-safe sound; proper retract attribution

F1: the per-expert boundary vote was a partial detector -- an inversion
inside a segment used by no other expert's boundary could escape, so
"re-route on any violation" overclaimed. build_gemma4_sorted_expert_tiles_bm32
now also runs a strided adjacent-pair scan: thread lid checks
indices[i-1] <= indices[i] for i = lid+1; i < M; i += 128, covering every
adjacent pair in [1, M) exactly once (1..128 iterations at the reachable
M in {4096,8192,16384}). Adjacent-pair monotonicity is transitive, so a
clean scan is a sound and complete sortedness oracle; it folds into the
same simd_or/threadgroup vote and the same retract (count[0]=0, count[1]=1).
The boundary checks stay as cheap, precise diagnostics.

F2: retracts were write-only in count[1] and surfaced as
fallback_metallib_unavailable -- misattribution in the only observable
surface. A dedicated fallback_sortedness_retracted counter now rides the
GemmA4 route counters and the C diagnostics ABI
(sizeof 80 -> 88, new uint64 at offset 80; existing offsets unchanged).
try_gemma4_expert_qmm returns the route class: count[0]==0 with count[1]==1
records fallback_sortedness_retracted, any other unusable build keeps
fallback_metallib_unavailable, then re-routes to the legacy path as before.

F4: new doctest drives the full armed() -> clear_and_arm() ->
snapshot_and_disarm() cycle and the attempts == hits + fallbacks invariant
including the new class; the route-table and counter-invariant tests now
cover fallback_sortedness_retracted.

Verified: cmake tests 262/262 + 3550 assertions pass; metal -Wall -Wextra
-fno-fast-math compile of kernels/quantized.metal is warning-free.

* perf(metal): E=256 expert-tile route + trust + gpu::eval UAF fix — darkbloom-base mirror (#7)

* perf(metal): instantiate E=256 expert-tile route for Qwen 3.5/3.6 MoE prefill (mirror of Cmlx/mlx 58fab46)

* fix(metal): use-after-free in gpu::eval for primitives that synchronize mid-eval (mirror)

* perf(metal): trust mode skips retract readback (mirror)

* fix(compile): preserve all-cache binding cleanup

---------

Co-authored-by: JasonHonKL <[email protected]>
Co-authored-by: AK <[email protected]>
Co-authored-by: Cheng <[email protected]>
Co-authored-by: Adityaj0 <[email protected]>
Co-authored-by: anchor <[email protected]>
Co-authored-by: codeAnqiang-ma <[email protected]>
Co-authored-by: Rohan Gautam <[email protected]>
Co-authored-by: Ayaan Gazali <[email protected]>
Co-authored-by: Erwin Zhang <[email protected]>
Co-authored-by: katlun-lgtm <[email protected]>
Co-authored-by: katlun-lgtm <[email protected]>
Co-authored-by: robertomeroni <[email protected]>
Co-authored-by: Fu Xiaonan <[email protected]>
Co-authored-by: Fu Xiaonan <[email protected]>
Co-authored-by: Hao Xu <[email protected]>
Co-authored-by: Feli <[email protected]>
Co-authored-by: Feli <[email protected]>
Co-authored-by: Eyüp Can Akman <[email protected]>
Co-authored-by: Cheng <[email protected]>
Co-authored-by: yentur <[email protected]>
Co-authored-by: Alessio Pollero <[email protected]>
Co-authored-by: Zhiqi Zhang <[email protected]>
Co-authored-by: Daniel Hiltgen <[email protected]>
Co-authored-by: Tanish Jain <[email protected]>
Co-authored-by: hojin12312 <[email protected]>
Co-authored-by: Xiang Chen <[email protected]>
Co-authored-by: x14ngch3n <[email protected]>
Co-authored-by: Duhyeon, Kim <[email protected]>
Co-authored-by: rohith <[email protected]>
Co-authored-by: Ishaan Samantray <[email protected]>
Co-authored-by: Aaishwarya Mishra <[email protected]>
Co-authored-by: Anastasiia Filippova <[email protected]>
Co-authored-by: Yanzhao Wang <[email protected]>
Co-authored-by: XXXXRT666 <[email protected]>
Co-authored-by: Jake Bowhay <[email protected]>
Co-authored-by: vraj patel <[email protected]>
Co-authored-by: Gusanidas <[email protected]>
Co-authored-by: Dwijen Patel <[email protected]>
Co-authored-by: Vladimir Iglovikov <[email protected]>
Co-authored-by: Brian C. <[email protected]>
Co-authored-by: Daniel Hiltgen <[email protected]>
Co-authored-by: YH Yan <[email protected]>
Co-authored-by: katlun-lgtm <[email protected]>
Co-authored-by: anupsv <[email protected]>
Co-authored-by: Gajesh Naik <[email protected]>
Co-authored-by: David Tai <[email protected]>

Upstream ml-explore/mlx c7ff35d ("Skip unnecessary simdgroup
computations for quantised MOE matmuls on NAX", ml-explore#4352) reached this fork
through the v0.32.2 squash (734241b), not by ancestry, in the form that
ml-explore#4572 later gave it: `const bool sg_active = sgp_sm > 0;` and two
`if (sg_active)` blocks around the MMA loop in each of
fp_gather_qmm_rhs_nax (fp_quantized_nax.h) and affine_gather_qmm_rhs_nax
(quantized_nax.h). This removes the flag and runs the MMA in every
simdgroup again, as before ml-explore#4352. It is the same change as the earlier
revert of ml-explore#4352 on 0.32.2 (e9debbe) and its re-application on the
0.32.3 kernels, hand-ported because the surrounding lines differ here.

Why: these kernels run only on GPUs with NAX. On a two-Mac JACCL pair
with one NAX GPU, mixture-of-experts serving with the skip present lost
all-reduce completions, and 92 to 96 GiB stayed wired on the host until
it was rebooted. The builds that served that pair afterwards carried the
revert. It is therefore kept out of every build that runs distributed on
mixed hardware.

Output is unchanged: stores stay bounded by sgp_sm, and the rows past
the tile that the MMA now reads are either inside M (rows_in_bounds) or
loaded as zeros by load_safe with sgp_sm rows.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

Set base.hash to the new merge-base. Remove the v0.32.2 sync section
and the gather_mm row-tile backport section, because upstream now has
those changes. Keep the revert of the ml-explore#4352 simdgroup skip in its own
section. Reduce the floor_divide section to the tests the fork keeps.
Add the re-deleted bypass-list workflow and scripts to Fork tooling.

Keep the revert of the ml-explore#4352 simdgroup skip in fp_gather_qmm_rhs_nax
and affine_gather_qmm_rhs_nax. The merge had taken upstream's
sg_active guard there.

State the new upstream sync policy in FORKDIFF.md: merge upstream
when David directs it, and move base.hash at each merge.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* Fix Log and Equal is_equivalent ignoring primitive state (ml-explore#4266)

Co-authored-by: Cheng <[email protected]>

* Stabilize reduced-precision InstanceNorm (ml-explore#4230)

* chore: Normalize negative axes in sort and argsort (ml-explore#4332)

* Clean up main thread compile cache before python interpreter shuts down (ml-explore#4373)

* chore: Check malformed jaccl hostfile that miss rdma in pairs (ml-explore#4284)

Co-authored-by: Cheng <[email protected]>

* Round mxfp8 block scales up to avoid saturation (ml-explore#4353)

Co-authored-by: Daniel Hiltgen <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Add support for the __array_namespace_info__  (ml-explore#4334)

* Stop a failed CUDA graph commit from poisoning the encoder (ml-explore#4356)

Co-authored-by: Cheng <[email protected]>

* Fix quantized kernels in JIT build (ml-explore#4372)

Co-authored-by: Cheng <[email protected]>

* Avoid zero work in stride-2 ConvTranspose3d (ml-explore#4343)

* [CUDA] Ce fused kernel (ml-explore#3947)

* Fix cpu exclusive scan for complex numbers (ml-explore#4272)

Co-authored-by: Cheng <[email protected]>

* Support Relocatable CUDA DLLs on Windows (ml-explore#4382)

* Use cast_to for fused AsType in compiled Metal kernels (ml-explore#4351)

Co-authored-by: katlun-lgtm <[email protected]>
Co-authored-by: Cheng <[email protected]>

* python: Declare DLPackCompatible protocol members as methods (ml-explore#4384)

* Fix quantizing sliced arrays (ml-explore#4381)

* Fix einsum dropping a trailing empty subscript (ml-explore#4299)

Co-authored-by: Cheng <[email protected]>

* Add script to run python tests (ml-explore#4393)

* Hold GIL in AttachedData destructor (ml-explore#4391)

* Patch bump to 0.32.3 (ml-explore#4406)

* Fix scan and sort for a zero-size axis (ml-explore#4340)

* python: Add a correction parameter in std and var (ml-explore#4348)

Signed-off-by: Pradyot Ranjan <[email protected]>
Co-authored-by: Pradyot Ranjan <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Fix tests/run.py stuck in macOS CI (ml-explore#4410)

* chore: Update extension example (ml-explore#4412)

* Fix sorted gather_qmm NAX row overflow above 32K (ml-explore#3922)

Co-authored-by: Cheng <[email protected]>

* Fix deadlock caused by mx.clear_streams() holding GIL (ml-explore#4413)

* Fix integer pow zeroing a whole SIMD vector on a negative exponent (ml-explore#4354)

* Make concurrency cap on load adaptative so that I/O scales with the machine (ml-explore#4408)

* Fix pad with an axes subset and negative axes (ml-explore#4364)

* Use sdpa_vector_2pass_1_gqa for GQA size 12 and 16 (ml-explore#4380)

* python: Make axis of put_along_axis optional (ml-explore#4360)

* python: Make axis default to -1 in take_along_axis (ml-explore#4368)

* Propagate NaN in arg reductions (ml-explore#4291)

Co-authored-by: Cheng <[email protected]>

* Fix mx.from_fp8 E4M3FN NaN decode with branchless carry (ml-explore#4376)

Co-authored-by: Saud <[email protected]>

* Update rules for PR limitation bypass list (ml-explore#4415)

* [CUDA] Use native events for GPU fence waits (ml-explore#4401)

Co-authored-by: Cheng <[email protected]>

* Validate GGUF tensor dimensions (ml-explore#4378)

Co-authored-by: Cheng <[email protected]>

* Fix crash on ellipsis indexing with too many trailing indices (ml-explore#4396)

Co-authored-by: Anastasiia Filippova <[email protected]>

* [CUDA] Cholesky via cuSOLVER (ml-explore#4208)

Co-authored-by: Cheng <[email protected]>

* python: Fix __setitem__ with negative index preceded by None (ml-explore#4414)

* Update nanobind to 3.0.1 (ml-explore#4417)

* chore: Use normalize_axis_index in scan ops (ml-explore#4383)

* python: Add explicit index and bytes support (ml-explore#4388)

* docs: fix typo recived -> received (ml-explore#4424)

Signed-off-by: Vaibhav Srivastava <[email protected]>

* Fix Metal grid sizing for strided scans (ml-explore#4420)

Co-authored-by: Cheng <[email protected]>

* Fix SGD and Adafactor weight decay mutating caller arrays (ml-explore#4426)

* Widen float16/bfloat16 to float32 in cpu reduce (ml-explore#4387)

* Fix sdpa_vector_2pass_1_gqa kernel batch offset for K/V in gqa decode (ml-explore#4431)

* Make head-dim-256 prefill with array mask run fused NAX kernel (ml-explore#4416)

* python: Make unstack return tuple instead of list (ml-explore#4448)

* python: Make iter() throw TypeError for 0-dim array (ml-explore#4425)

Co-authored-by: Cheng <[email protected]>

* Fix save and save_safetensors corrupting a lazily loaded source file (ml-explore#4434)

* [Metal] Use hypot for complex ops (ml-explore#4345)

Co-authored-by: Cheng <[email protected]>

* Add get_array_buffer_size to query buffer size for arrays (ml-explore#4436)

* Add M5 ultra tunings for non-quantized matmuls (ml-explore#4447)

* [CUDA] Fix completion worker busy loop (ml-explore#4452)

* Fix non transposed affine qmm dispatch logic (ml-explore#4392)

* Fix sorted gather_qmm on ragged K (ml-explore#4009)

Co-authored-by: Cheng <[email protected]>

* Don't restore a thread-affine stream when StreamContext dies on another thread (ml-explore#4462)

* Fix Pad vjp with an axes subset and negative axes (ml-explore#4441)

Co-authored-by: Cheng <[email protected]>

* Fix integer div-by-zero crash in CPU backend (sort, scan, binary ops) (ml-explore#4442)

Co-authored-by: Cheng <[email protected]>

* Use 64-bit file seeks on Windows (ml-explore#4456)

Co-authored-by: Cheng <[email protected]>

* Use precise::exp in Sigmoid so compiled and eager sigmoid agree (ml-explore#4461)

Co-authored-by: Pierre Lamy <[email protected]>

* Use NAX attention for long unmasked D72/D80 inputs (ml-explore#4455)

* Add D512 support to Metal vector attention (ml-explore#4459)

Co-authored-by: Cheng <[email protected]>

* Add a scoped_env utility for tests (ml-explore#4474)

* [Metal] global scale for qmm (ml-explore#4458)

* Fix re-entering the same stream context manager (ml-explore#4478)

* Leak global CommandEncoder on exit (ml-explore#4480)

* Update CONTRIBUTING.md (ml-explore#4469)

* python: Fix bytes() on non-contiguous arrays (ml-explore#4449)

* Break the sibling cycle when an array is released by assignment (ml-explore#4453)

Co-authored-by: Cheng <[email protected]>

* win: Use DXGI for accurate WDDM VRAM memory budget (ml-explore#4457)

Co-authored-by: Cheng <[email protected]>

* Add type hints and docstrings for ALiBi layer (ml-explore#4472)

* Fix fp quantized matmul corruption when the quantized dim is not a multiple of 32 (ml-explore#3912)

Co-authored-by: Rohan Gautam <[email protected]>

* Fix CUDA test synchronization flake (ml-explore#4490)

* ci: Skip bypass list updates in forks (ml-explore#4495)

* Load global scales in qmm_t kernels (ml-explore#4483)

* python: Release the GIL when copying Metal DLPack inputs (ml-explore#4497)

Signed-off-by: Ranran Haoran Zhang <[email protected]>

* Use matrix kernels for global-scale gather_qqmm (ml-explore#4481)

* Adding metal kernels for the gated delta nets. (ml-explore#4020)

Co-authored-by: Yanzhao Wang <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Fix jaccl ring all_gather: direction 1 slice was not mirrored (ml-explore#4443)

Co-authored-by: Dave Hind <[email protected]>
Co-authored-by: Anastasiia Filippova <[email protected]>

* Skip Metal-only gated delta kernel tests on other backends(fix CI) (ml-explore#4522)

Co-authored-by: Anastasiia Filippova <[email protected]>

* [CUDA] Add global scale support to gather_qmm (ml-explore#4507)

Co-authored-by: Anastasiia Filippova <[email protected]>

* [BUG][Metal] Deadlock in fence (ml-explore#4552)

Co-authored-by: Cheng <[email protected]>

* Fix ordering of complex64 with NaNs (ml-explore#4519)

* Improve metal memory usage for SDPA D256 (ml-explore#4505)

Co-authored-by: Cheng <[email protected]>

* Add support for matrix transpose (ml-explore#4402)

Co-authored-by: Cheng <[email protected]>

* Improve metal memory usage for SDPA D512 (ml-explore#4487)

* Support shapeless compilation of scan operations (ml-explore#4510)

* [Metal] Gather mm improvement (ml-explore#4567)

* Fix perf regression for upsample `linear` mode align_corners=False (ml-explore#4500)

* Add fused Metal kernels for fast.cross_entropy (ml-explore#4520)

Co-authored-by: Anastasiia Filippova <[email protected]>

* python: Add missing defaults to tri, tril, triu, gather_mm signatures (ml-explore#4492)

* Add Metal SDPA support for D96/V64 (ml-explore#4499)

Co-authored-by: Cheng <[email protected]>

* Raise IndexError for out of bounds axes (ml-explore#4484)

* Fix floor_divide for integers (ml-explore#4515)

Co-authored-by: Cheng <[email protected]>

* Fix cpu binary ops on large data (ml-explore#4517)

Co-authored-by: Cheng <[email protected]>

* Do not throw from CUDA destructors and avoid implicit default streams in eval/compile (ml-explore#4514)

Co-authored-by: Cheng <[email protected]>

* [Metal] gather_qmm improvement  (ml-explore#4572)

* Fix floating-point constant precision in compiled kernels (ml-explore#4511)

* chore: Fix assertion in ReduceScatter::eval_cpu (ml-explore#4557)

* Leak thread local streams on exit in main thread (ml-explore#4576)

* Adding VJP for the scaled dot product attention. (ml-explore#4563)

* python: Add __pos__ support to array (ml-explore#4488)

Co-authored-by: Cheng <[email protected]>

* Add Python 3.15 and 3.15t support (ml-explore#4555)

Co-authored-by: Cheng <[email protected]>

* Add fixed size `unique` operation  (ml-explore#4501)

Co-authored-by: Cheng <[email protected]>

* Fix normalizing single-element slice with negative stride (ml-explore#4446)

Co-authored-by: Cheng <[email protected]>

* Patch bump to 0.32.4 (ml-explore#4582)

* Adding VJP for gated delta net. (ml-explore#4565)

* Fix `col_reduce_longcolumn` for negative-stride views (ml-explore#4529)

Co-authored-by: Cheng <[email protected]>

* Clear CUDA errors once reported and raise eval errors in buffer protocol (ml-explore#4523)

Co-authored-by: Cheng <[email protected]>

* Improve pre-M5 metal memory usage for SDPA D512 (ml-explore#4518)

* Fix split data size span (ml-explore#4540)

* Re-introduce Array API Tests suite (ml-explore#4433)

Signed-off-by: Pradyot Ranjan <[email protected]>
Co-authored-by: Pradyot Ranjan <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Fix view inside of a vmap (ml-explore#4586)

* Fix cumsum of a non-contiguous bool (ml-explore#4590)

* docs: contraction dtype requirements in matmul/einsum (ml-explore#4550)

Co-authored-by: Cheng <[email protected]>

* Fixing cooperative tensor build issues on Mac OS 27 (ml-explore#4594)

* Handle EOF and EINTR in ParallelFileReader::read (ml-explore#4532)

Co-authored-by: Cheng <[email protected]>

* Clamp slice bounds when a slice is combined with an array index (ml-explore#4513)

Co-authored-by: Cheng <[email protected]>

* Fix grouped Conv1d/Conv2d init to use in_channels // groups (ml-explore#4544)

Co-authored-by: Cheng <[email protected]>

* Fix mx.view misaligned offset (ml-explore#4585)

Co-authored-by: Cheng <[email protected]>

* [CUDA] Add fp_gather_qmv to optimize gather_qmm (ml-explore#4508)

* Fix eval cleanup deadlock by committing GPU streams before waiting (ml-explore#4556)

* Fix shared_buffer_reshape contiguity flags (ml-explore#4589)

* Fix view last axis stride (ml-explore#4587)

Co-authored-by: Cheng <[email protected]>

* Fix empty unsigned sum/prod failing on Metal (ml-explore#4546)

Co-authored-by: Cheng <[email protected]>

* [CUDA] Add gather_qmm_rhs_sm80 (ml-explore#4554)

Co-authored-by: Cheng <[email protected]>

* python: Fix bare ellipsis __setitem__ by squeezing leading singleton dimensions (ml-explore#4561)

* Fix row/col contiguous flags in view (ml-explore#4588)

* Do not cache failed distributed backend initialization (ml-explore#4558)

* Fix reduction over large arrays (ml-explore#4549)

Co-authored-by: Cheng <[email protected]>

* docs: Correct the export callback type in the guide (ml-explore#4578)

* Fix MultiOptimizer crash when an optimizer's group is empty (ml-explore#4570)

* Fix fft size one (ml-explore#4591)

* Sanitize '-' in make_template_hash for metal kernel (ml-explore#4584)

Signed-off-by: GokayAI <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Fix as_strided contiguity (ml-explore#4575)

* Reject non-aligned pointers when converting to mx.array (ml-explore#4592)

* Fix conjugation in complex VJPs for inverse trig ops, cos, sqrt, and rsqrt (ml-explore#4606)

* Fix complex to bool (ml-explore#4593)

* [jaccl] second cable for the mesh (ml-explore#4607)

* Fix vmap of scatter placing the batch axis wrongly in the updates (ml-explore#4322)

Co-authored-by: Cheng <[email protected]>

* Unroll TK=4 keys per iter in sdpa_vector_2pass_1 (ml-explore#4596)

Co-authored-by: Cheng <[email protected]>

* Fix CPU float16 SIMD comparison masks (ml-explore#4602)

Co-authored-by: Cheng <[email protected]>

* Add jvp support for logcumsumexp (ml-explore#4611)

* [CUDA] Export CommandEncoder and Device (ml-explore#4619)

* Use the fast qmv kernel for outputs not divisible by 8 (ml-explore#4516)

Co-authored-by: Cheng <[email protected]>

* Fix ResidencySets reading autoreleased error object (ml-explore#4623)

* Use the small-channel implicit gemm loader for C_per_group = 8 (ml-explore#4617)

* [CUDA] Reduce floor for fp_gather_qmv (ml-explore#4620)

* Fix default.metallib lookup under swift test with Swift Build (ml-explore#4527)

* Fix metal validation error when passing empty vector (ml-explore#4622)

Co-authored-by: Cheng <[email protected]>

* Replace Black formatter with ruff (ml-explore#4600)

* Tune M1 quantized matmul dispatch for six-row large shapes (ml-explore#4629)

Co-authored-by: deadbydawn <[email protected]>
Co-authored-by: Cheng <[email protected]>

* [metal] Add env var based residency refresh interval (ml-explore#4633)

* Promote integer inputs in logcumsumexp (ml-explore#4625)

Co-authored-by: Cheng <[email protected]>

* Fix slicing's non contiguous start (ml-explore#4599)

* Fix float16 overflow in clip_grad_norm (ml-explore#4626)

* Collapse the non sorted axes to pick the contiguous sort kernel (ml-explore#4366)

Co-authored-by: Cheng <[email protected]>

* Fix 32 bit dispatch grid overflow in Metal reductions (ml-explore#4437)

Co-authored-by: Cheng <[email protected]>

* Fix gpu boolean masking in 0 dim array (ml-explore#4635)

Co-authored-by: Cheng <[email protected]>

* Fix vmap of an inverse real FFT to an odd length (ml-explore#4637)

* Fix cross-process race in CPU compile cache (ml-explore#4638)

Co-authored-by: Cheng <[email protected]>

* Update acknowledgments for past pinv contribution (ml-explore#4645)

* Fix foreign buffer ownership (ml-explore#4634)

* Add MLX_CONV_WINOGRAD switch to disable winograd for conv2d (ml-explore#4639)

* Keep split-K quantized matmul partials in float32 (ml-explore#4641)

* Thin N GEMM kernels (ml-explore#4640)

Co-authored-by: Cheng <[email protected]>

* Add 1-bit affine quantization support (Metal) (ml-explore#3161)

Co-authored-by: Angelos Katharopoulos <[email protected]>
Co-authored-by: bri-prism <[email protected]>

* Update Metal Complex division in numpy style to avoid overflow condition (ml-explore#4644)

* Fix pickle bfloat16 strides (ml-explore#4649)

* Fix metal negative stride index in masked_scatter/scatter/SDPA mask/segmented_gemm (ml-explore#4650)

* Fix missing break binary bool (ml-explore#4651)

* Fix cpu binary_op_dispatch_dims on large data (ml-explore#4648)

* Add FloorDivide primitive for integer floor division (ml-explore#4642)

Co-authored-by: Cheng <[email protected]>

* [Metal] Fix SDPA max threadgroup for M1/M2 generation (ml-explore#4643)

Co-authored-by: Cheng <[email protected]>

* Fix incorrect output mask computation in `mx.block_masked_mm` (ml-explore#4652)

* Add thin M GEMM dispatch (ml-explore#4654)

* Fix Python dtpye -> mlx conversion narrowing (ml-explore#4656)

* Fix input/output scalar merging in compile (ml-explore#4658)

Co-authored-by: Cheng <[email protected]>

* Validate .npy shape in mlx/io/load.cpp (ml-explore#4657)

Co-authored-by: Cheng <[email protected]>

* Replace "pull_request_target" event with "pull_request" in update_bypass_list workflow (ml-explore#4667)

* [CUDA] Only count consecutive misses in the LRUCache thrashing check (ml-explore#4661)

Co-authored-by: Cheng <[email protected]>

* Fix conv input gradient shape with uneven or large padding (ml-explore#4662)

* Fix missing stream in layer_norm and sdpa fallbacks (ml-explore#4663)

Co-authored-by: Cheng <[email protected]>

* Fix direct CPU conv writing past output (ml-explore#4668)

Co-authored-by: Cheng <[email protected]>

* Add stft and istft to mlx.core.fft (ml-explore#3639)

* Fix lost rank output in the distributed launcher (ml-explore#4615)

Co-authored-by: Cheng <[email protected]>

* Fix split-K test compatibility on M5, M6 (ml-explore#4677)

* Fix signedness warning in quantized Metal tail bounds check (ml-explore#4675)

* Fix batched row count mismatch in metal qvm_split_k (ml-explore#4673)

Co-authored-by: Cheng <[email protected]>

* Fix metal conv bounds checks and integer overflows (ml-explore#4672)

Co-authored-by: Cheng <[email protected]>

* python: Make arctan2/radians/degrees support scalar inputs (ml-explore#4679)

Co-authored-by: Cheng <[email protected]>

* Fix missing add_temporary in pad_gpu/fill_gpu (ml-explore#4660)

Co-authored-by: “Ronald <“[email protected]”>

* Enable non-transposed NAX gather_qmm_rhs (ml-explore#4684)

* Add Metal extension build support (ml-explore#4004)

* Adding minmax method to MLX to follow Numpy/Jax API (ml-explore#4674)

Co-authored-by: Cheng <[email protected]>

* Fix missing stream in gather matmul indices cast (ml-explore#4687)

* [BUG] Metal fence deadlock (ml-explore#4689)

* Patch bump to 0.32.5 (ml-explore#4707)

* docs: Publishing a New Release (ml-explore#4691)

* ci: Run update_bypass_list on push instead of pull_request (ml-explore#4671)

* [CUDA] Run the vector SDPA kernels at head_dim 256 (ml-explore#4698)

Co-authored-by: yanghaojin <[email protected]>

* [CUDA] Fix grouped gemm with large data (ml-explore#4696)

Co-authored-by: yanghaojin <[email protected]>
Co-authored-by: Cheng <[email protected]>

* Fix conv2d OOB reads for partial output-channel tiles (ml-explore#4653)

* docs: MetalLibrary in custom extension guide (ml-explore#4701)

* [Metal] Fix SDPA f16/f32 d512 max threadgroup for M1/M2 generation (ml-explore#4708)

* Add float16 head_dim 256 TK=1 SDPA vector kernel for M1/M2 (ml-explore#4719)

Co-authored-by: Cheng <[email protected]>

* Allow configuring the CPU JIT cache directory (ml-explore#4714)

Co-authored-by: Cheng <[email protected]>

* Make `broadcast_arrays` return a tuple (ml-explore#4713)

Co-authored-by: Cheng <[email protected]>

* Fix unaligned scalar reads and oversized contiguous buffer views (ml-explore#4721)

Co-authored-by: Cheng <[email protected]>

* [Metal] Add staged affine_qmm_t_nax kernels with split-k (ml-explore#4682)

Signed-off-by: Yanzhao Wang <[email protected]>
Co-authored-by: Cheng <[email protected]>

* [Metal] Add fast route for dequantizing fp4 to half in fp_qmm_t (ml-explore#4709)

Co-authored-by: Cheng <[email protected]>

* [Metal] Add qmv_mma/qmv_nax kernels for 6-16 rows with large N/K (ml-explore#4712)

Co-authored-by: Cheng <[email protected]>

* docs(forkdiff): move the base to upstream c6abf0d after the merge

Set base.hash to the new merge-base. Remove the v0.32.2 sync section
and the gather_mm row-tile backport section, because upstream now has
those changes. Keep the revert of the ml-explore#4352 simdgroup skip in its own
section. Reduce the floor_divide section to the tests the fork keeps.
Add the re-deleted bypass-list workflow and scripts to Fork tooling.

Keep the revert of the ml-explore#4352 simdgroup skip in fp_gather_qmm_rhs_nax
and affine_gather_qmm_rhs_nax. The merge had taken upstream's
sg_active guard there.

State the new upstream sync policy in FORKDIFF.md: merge upstream
when David directs it, and move base.hash at each merge.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix: repair the merge in the quantized kernels and the lint

The merge put upstream's new "typename U = T" template parameter on
the fork's qmm_t_expert_impl instead of qmm_t_impl. Move it back.

fp_gather_qmv_fast_tail now calls adjust_matrix_offsets with the
global scale, as upstream's fp_gather_qmv_fast does.

Remove the extra closing brace that the previous commit put in
affine_gather_qmm_rhs_nax. The kernel body is again the fork's body.

Import numpy in test_fast_sdpa.py again. Upstream removed the import,
and the fork's FP32-partials tests use it. Wrap the long lines in
check_forkdiff.py for the ruff E501 rule that upstream turned on.

Remove the PR #15 bias-sum sentence from fork.yaml: upstream now
promotes the bias-sum operands itself (ml-explore#4516).

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* test: follow the upstream allocator and gather_qmm APIs in fork tests

Upstream allocator::malloc now returns an owning allocator::Data, and
allocator::free is gone. The fork's snapshot and footprint tests need
explicit release points, so they use allocator::allocator().malloc and
free, as upstream's own allocator tests do.

Upstream gather_qmm takes a global_scale argument before
sorted_indices. Pass std::nullopt in the GPT-OSS MXFP4 tests.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix(cpu): floor_divide of min by -1 must not trap

Upstream's FloorDivide divides the minimum signed value by -1 on the
CPU. The quotient overflows, and on x86 the division traps (SIGFPE).
The fork test "test floor_divide integer edge cases on the CPU"
crashed in the Fedora x86_64 job.

Divide by 1 there instead, as the PR #32 functors did. The result
wraps to min. Describe the guard in the floor_divide section of
fork.yaml.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* Compute the JACCL ring wire offsets in 64 bits

CodeQL (cpp/integer-multiplication-cast-to-long) flagged the ring's wire
offset expressions on PR 38: `lr * n_wires` is an `int` product that is
widened to 64 bits only after the multiplication. The same pattern is in
three ring_impl.h loops and in the two `size_ * 2 * n_conns_` thresholds in
ring.cpp. Widen the first operand to int64_t so every product is computed
in 64 bits. The operands are small, so results do not change. fork.yaml
describes the new fork change.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* Compute the vector SDPA query layout check in 64 bits

CodeQL (cpp/integer-multiplication-cast-to-long, alert 29) fails the PR 38
check on the vector SDPA query layout test: `shape[3] * shape[bidx]` is an
`int` product compared with a 64-bit stride. Widen `shape[3]` to int64_t.
The shapes are small, so the result does not change. The fork.yaml section
for the JACCL ring fix now covers this file too.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* Fix `mx.take`'s negative indices (ml-explore#4729)

* Fix normal loc scale shape (ml-explore#4730)

* Fix zero-copy numpy arrays dtype mismatch (ml-explore#4731)

* ci: skip the 4 GiB reduce test where the runner cannot hold it

Upstream's test_init_reduce_large_output (ml-explore#4437) makes a 4 GiB boolean
output. The fork CI runs on a GitHub-hosted macOS runner with a
3.5 GiB Metal buffer limit and 7 GB of memory. On the GPU the
allocation fails; on the CPU it competes with the other test files
for memory. Skip the test when the default device cannot hold the
output. Machines with more memory still run it.

When a CPU test file fails, print the end of its log at once. The job
of the previous run was stopped before the full logs printed, so the
cause of the test_eval failure on the CPU is not known.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* Widen affine qmv_wide row tiles to 8 (ml-explore#4728)

* Fix/padding jvp vjp (ml-explore#4736)

* fix: qqmm refuses a global scale on the CPU; adapt GPU-path tests

The CPU QQMatmul kernel has no global scale input and ignores the
scales that qqmm passes to it. qqmm with global scales on a CPU stream
returned a wrong result with no error: in the fork CPU test pass,
test_fp_qmv_fast_unaligned_n_global_scale differed by about 1000.
qqmm now raises "Global scale is only supported on the GPU" on a
stream that is not a GPU stream, as gather_qmm already does.

Upstream tests that check only GPU kernels failed in the CPU pass:
test_sdpa_vector_2pass_threadgroup_limit (force_fused=True),
test_fp_qmv_fast_unaligned_n_global_scale,
test_qqmm_global_scale_matrix_kernels and
test_gather_qmm_global_scale. When the default device is the CPU,
they now check that the documented error is raised. The GPU pass runs
them as before.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* ci: add a temporary diagnostic for the CPU NaN; reword the 4 GiB skip

Add one temporary step after the build in metal-build-and-test. It
runs the failing case of test_fp_gather_qmm_non_multiple_of_32
(K=512, N=1040, transpose=False, sorted, float32) on the CPU with a
fresh allocator, after a warm-up that fills the allocator cache with
NaN buffers, and with the buffer cache off, and on the GPU. It prints
NaN counts per expert for w, w_hat, y_q and y_hat. The step has
continue-on-error, so the full test suite still runs. The next push
removes it.

State in fork.yaml when test_init_reduce_large_output skips, why it
cannot be smaller (ml-explore#4437), and that David ruled on 2026-10-23 to keep
the skip.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* ci: move the CPU NaN diagnostic to its own short job

Remove the temporary diagnostic step from metal-build-and-test. Its
reimplementation found no NaN in any mode.

Add a temporary job, diag-cpu-nan, that builds the wheel and runs only
the diagnostic, so its log is ready after the build. It runs the real
test_fp_gather_qmm_non_multiple_of_32 alone, then after the tests that
come before it, and an instrumented copy of the test body (only the
failing case, and all cases in the test order, each evaluated as one
graph). The findings also go to ::notice annotations. The next push
removes the job.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* ci: install numpy for the temporary NaN diagnostic job

The tests import numpy through mlx_tests. The diagnostic job installed
only the wheel, so every run stopped at the import.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* Fix data race in distributed ring on pending send/recv lists (ml-explore#4737)

Co-authored-by: Cheng <[email protected]>

* Use data_size for gather/scatter large-offset checks (ml-explore#4724)

Co-authored-by: Cheng <[email protected]>

* fix(cpu): clear C before the BLAS matmul when beta is 0

test_fp_gather_qmm_non_multiple_of_32 (K=512, N=1040, sorted,
float32) gave NaN on the CPU. The diagnostic job found the NaN only in
the reference y_hat = x @ w_sel, a batched matmul with M = 1. Its
inputs had no NaN. The test failed alone too, but only after the
earlier subtests of the same test: y_hat is 266240 bytes, the same
size as the packed fp4 weights of the previous subtest, so the output
reused that buffer from the allocator cache. Accelerate read C even
with beta = 0 and computed NaN * 0.

Fill C with zeros before cblas_sgemm, cblas_dgemm and cblas_cgemm
when beta is 0. MLX's own SIMD gemm already skips C when beta is 0.

Remove the temporary diagnostic job.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* fix(cpu): clear C at every BLAS gemm call with beta 0

The previous commit cleared C only in the gemms/cblas.cpp wrappers.
Four more call sites call cblas_sgemm directly with beta 0:
BlockMaskedMM and GatherMM (the CPU MoE gathered matmul) in
masked_mm.cpp, and the two explicit-gemm convolutions in conv.cpp.
They can give the same NaN * 0 from a reused output buffer.

Add one helper, clear_gemm_output in gemm.h, that fills the exact C
region of a call (rows by columns, with its ldc) with zeros. Use it at
all five call sites, the wrappers included. No other BLAS gemm or gemv
call exists in the CPU backend.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* ci: build once and run the Metal and CPU test passes in parallel

Split the fork-tests job into parallel jobs. build-wheel and build-cpp
build in parallel. build-cpp also runs the C++ tests on the GPU and on
the CPU. python-gpu (3 shards) and python-cpu (4 shards) run the Python
tests from the uploaded wheel. metal-build-and-test merges the coverage
profiles of all jobs, writes the same coverage report, and fails when
another job failed.

python_shards.py checks the shard map: each python/tests/test*.py file
must be in each pass exactly one time. Split files must partition their
test IDs. The last job fails when a test ran in two shards.

fork-build-cache restores and saves ccache on pull requests and on main,
and keeps a cache of the compiled Metal kernels through an xcrun wrapper
keyed on the preprocessed source, the arguments and the toolchain.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* [CUDA] Fix gather_qmm for experts past 2^32 weight elements (ml-explore#4738)

Co-authored-by: Cheng <[email protected]>

* ci: split the long CPU test files and skip cache saves with no misses

Split test_fast_sdpa and test_quantized into named tests and rest
entries in the CPU shards, so that no CPU shard has an entry longer than
test_gated_delta_sequential. Put the entries in longest-first order.

Save the ccache files and the Metal kernel cache only when the build had
misses.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* ci: key the build caches on compiler versions, not the runner image

The full --version text of clang and metal names the install folder and
the OS, and it changed between three macos-26 runner images, so every
job started from an empty cache. Key ccache on the clang major.minor
version and fall back to a cache of another clang version; ccache keeps
compiler_check=content. Key the Metal kernel cache on the metal compiler
version and the kernels hash, and fall back only to a cache of the same
metal compiler. Each Metal entry key holds the first line of the metal
--version text and the SDK build.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* Add array::is_uniquely_referenced for the fork C API

Upstream 5dbe999 (ml-explore#4634) makes is_donatable() also require that MLX owns
the buffer, so the allocator never reuses a wrapped external buffer. The
fork's mlx-c buffer-info API reported `unique` from is_donatable(), and
mlx-swift-lm uses that as a "no other reference" proof for paged KV
segments, checkpoint memory accounting and MTP state, including wrapped
host pages, which now always read as not unique.

Add is_uniquely_referenced(): the same reference-count test without the
ownership test. is_donatable() and donation behavior do not change.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* fix: async_eval must not throw a stream task error

test_eval.test_async_eval_error_in_synchronize (upstream ml-explore#4338) failed
on some runs of the fork's CPU test pass: mx.async_eval raised the
LU factorization error that the test expects only from
mx.synchronize.

Cause: StreamThread::enqueue throws the stream's stored error when it
runs on the main thread (ml-explore#3742). async_eval enqueues several tasks.
When the stream thread finishes the failing task before async_eval
enqueues the next one, that enqueue throws. The order depends on
thread timing, so the failure is intermittent. The fork does not
change this path.

async_eval now holds a scheduler::DeferStreamErrors guard. While it
lives, enqueue on that thread does not throw the stored error, and
synchronize throws it later. eval does not change.

Add a C++ test that stores the error before async_eval, so it fails
on every run without the fix.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

* Fix GPU pad for outputs larger than 2^31 elements (ml-explore#4744)

* ci: format python_shards.py with the pinned ruff

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* ci: run the Python tests in 2 GPU and 3 CPU shards

The hosted macOS runners are a shared pool. Pack the GPU pass into 2
shards and the CPU pass into 3 shards, balanced by the durations that
run 39902948153 measured. A run now uses 8 macOS jobs, and at most 6 run
at the same time (before: 10 jobs, 8 at the same time).

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* ci: give the NAX kernels Metal cache hits and lighten CPU shard 0

The preprocessed source of the NAX kernels names the Metal toolchain
mount folder, which has a random suffix on each runner. The 10 NAX
kernels missed the Metal cache in every warm build (41 of 51 hits).
Remove the suffix from the text that the key hashes, and keep the
toolchain version.

CPU shard 0 was the slowest shard in the cold and warm runs of b72420d
(651 s and 470 s of tests). Move test_quantized:rest to CPU shard 1
(221 s).

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>

* docs(forkdiff): move the base to upstream afc03c5

Set base.hash to the new merge-base after the merge of the 9 upstream
commits. No fork section changes: the merge has no conflict, and each
merged file gets exactly upstream's change.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

---------

Signed-off-by: Pradyot Ranjan <[email protected]>
Signed-off-by: Vaibhav Srivastava <[email protected]>
Signed-off-by: Ranran Haoran Zhang <[email protected]>
Signed-off-by: GokayAI <[email protected]>
Signed-off-by: Yanzhao Wang <[email protected]>
Co-authored-by: rohith <[email protected]>
Co-authored-by: Cheng <[email protected]>
Co-authored-by: Vladimir Iglovikov <[email protected]>
Co-authored-by: Brian C. <[email protected]>
Co-authored-by: Cheng <[email protected]>
Co-authored-by: Erwin Zhang <[email protected]>
Co-authored-by: Daniel Hiltgen <[email protected]>
Co-authored-by: Daniel Hiltgen <[email protected]>
Co-authored-by: Aaishwarya Mishra <[email protected]>
Co-authored-by: YH Yan <[email protected]>
Co-authored-by: Dwijen Patel <[email protected]>
Co-authored-by: Anastasiia Filippova <[email protected]>
Co-authored-by: Ayaan Gazali <[email protected]>
Co-authored-by: katlun-lgtm <[email protected]>
Co-authored-by: katlun-lgtm <[email protected]>
Co-authored-by: Adityaj0 <[email protected]>
Co-authored-by: Eyüp Can Akman <[email protected]>
Co-authored-by: Pradyot Ranjan <[email protected]>
Co-authored-by: Pradyot Ranjan <[email protected]>
Co-authored-by: Jagrit Digani <[email protected]>
Co-authored-by: Philip John Basile <[email protected]>
Co-authored-by: Alessio Pollero <[email protected]>
Co-authored-by: Duhyeon, Kim <[email protected]>
Co-authored-by: Atirna <[email protected]>
Co-authored-by: Saud Kamran <[email protected]>
Co-authored-by: Saud <[email protected]>
Co-authored-by: YH Yan <[email protected]>
Co-authored-by: Roshan Sharma <[email protected]>
Co-authored-by: sashko-zakharchuk <[email protected]>
Co-authored-by: XXXXRT666 <[email protected]>
Co-authored-by: Vaibhav Srivastava <[email protected]>
Co-authored-by: Robin Gieseke <[email protected]>
Co-authored-by: 2sumtech <[email protected]>
Co-authored-by: Vedant Singhal <[email protected]>
Co-authored-by: JasonHonKL <[email protected]>
Co-authored-by: Simeet Nayan <[email protected]>
Co-authored-by: Cdimoy <[email protected]>
Co-authored-by: Alessio Pollero <[email protected]>
Co-authored-by: Rohan Gautam <[email protected]>
Co-authored-by: michalk8 <[email protected]>
Co-authored-by: Nilesh Patil <[email protected]>
Co-authored-by: Marcos Ashton <[email protected]>
Co-authored-by: pierre427 <[email protected]>
Co-authored-by: Pierre Lamy <[email protected]>
Co-authored-by: Yanzhao Wang <[email protected]>
Co-authored-by: David Koski <[email protected]>
Co-authored-by: Isaac Hernández <[email protected]>
Co-authored-by: Tudorica Constantin-Alexandru <[email protected]>
Co-authored-by: Ritaban Mitra <[email protected]>
Co-authored-by: Ranran <[email protected]>
Co-authored-by: Tommaso Pegolotti <[email protected]>
Co-authored-by: Drifter4242 <[email protected]>
Co-authored-by: Dave Hind <[email protected]>
Co-authored-by: Valentin Roussellet <[email protected]>
Co-authored-by: Ken Wang <[email protected]>
Co-authored-by: Spass <[email protected]>
Co-authored-by: zsun6 <[email protected]>
Co-authored-by: rishabh <[email protected]>
Co-authored-by: Yanzhao Wang <[email protected]>
Co-authored-by: devangpratap <[email protected]>
Co-authored-by: Jasim Kareem <[email protected]>
Co-authored-by: Ryan Choi <[email protected]>
Co-authored-by: Ronald Mannak <[email protected]>
Co-authored-by: Mateusz Rajski <[email protected]>
Co-authored-by: Ishaan Potle <[email protected]>
Co-authored-by: Moinuddin Shaik <[email protected]>
Co-authored-by: Moinuddin Shaik <[email protected]>
Co-authored-by: shaeantoine <[email protected]>
Co-authored-by: Nicoló Angileri <[email protected]>
Co-authored-by: Sam Small <[email protected]>
Co-authored-by: Ishti <[email protected]>
Co-authored-by: GokayAI <[email protected]>
Co-authored-by: Valeriy Sofin <[email protected]>
Co-authored-by: Yash Raj Pandey <[email protected]>
Co-authored-by: Kaiyu Shi <[email protected]>
Co-authored-by: crabel99 <[email protected]>
Co-authored-by: nav <[email protected]>
Co-authored-by: LongYinan <[email protected]>
Co-authored-by: Joseph Heck <[email protected]>
Co-authored-by: deadbydawn101 <[email protected]>
Co-authored-by: deadbydawn <[email protected]>
Co-authored-by: Vijay Kapse <[email protected]>
Co-authored-by: Long Yixing <[email protected]>
Co-authored-by: Aditya Dhulipala <[email protected]>
Co-authored-by: Ishaan Mehta <[email protected]>
Co-authored-by: Mattt <[email protected]>
Co-authored-by: Irakli Salia <[email protected]>
Co-authored-by: Pasha Khosravi <[email protected]>
Co-authored-by: Angelos Katharopoulos <[email protected]>
Co-authored-by: bri-prism <[email protected]>
Co-authored-by: CodeAndCanvas728 <[email protected]>
Co-authored-by: SounLabs <[email protected]>
Co-authored-by: David Dalcu <[email protected]>
Co-authored-by: Akshay Hitesh Shah <[email protected]>
Co-authored-by: Sash <[email protected]>
Co-authored-by: tc3oliver <[email protected]>
Co-authored-by: Magnus Lundstedt <[email protected]>
Co-authored-by: jrp2014 <[email protected]>
Co-authored-by: Yi <[email protected]>
Co-authored-by: “Ronald <“[email protected]”>
Co-authored-by: Ivan Petrukha <[email protected]>
Co-authored-by: Haojin Yang <[email protected]>
Co-authored-by: yanghaojin <[email protected]>
Co-authored-by: Marcin Markiewicz <[email protected]>
Co-authored-by: Corentin Kerisit <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
Co-authored-by: Ambisphaeric <[email protected]>
Co-authored-by: Sankur Kundu <[email protected]>
Co-authored-by: begcdn <[email protected]>
Co-authored-by: Core <[email protected]>
Co-authored-by: David Tivris <[email protected]>

Brings in main at f78f0fa (2026-10-24): the upstream ml-explore/mlx sync
to afc03c5 (#38, with its history restored by #42) and the fork fixes
#40 and #43-#46.

Conflicts and how they were resolved:
- mlx/distributed/jaccl/lib/jaccl/mesh.cpp: main's per-wire queue setup
  (upstream ml-explore#4607, a second cable for the mesh) with this branch's
  COMPLETION_QUEUE_DEPTH for every completion queue.
- mlx/distributed/jaccl/lib/jaccl/mesh_impl.h: main's stride parameters
  for all_gather and sum_scatter, with this branch's progress-guard calls
  and stage_send in every send path.
- mlx/distributed/jaccl/lib/jaccl/ring_impl.h: this branch keeps its
  guarded member dispatch_wires (pool-wire errors are reported, the group
  closes after all wires return). Main moved an unguarded copy into
  threadpool.h and called it as dispatch_wires(pool_, ...); the ring call
  sites call the member again. The all_gather slice code is main's
  (identical to this branch's backport of upstream ml-explore#4443).
- quantized_nax.h, fp_quantized_nax.h: main's kernels. Main already keeps
  the revert of the upstream ml-explore#4352 simdgroup skip (no sg_active) and adds
  the ragged-K tail bound of upstream ml-explore#4009.
- fork.yaml: main's file, plus the JACCL description and section. The
  upstream backports this branch carried (ml-explore#4443, ml-explore#4557, ml-explore#4558, the
  scoped_env helper of ml-explore#4474) are now upstream ancestry, so their "not fork
  work" section and its references are gone. The JACCL section now says
  that a mesh with more than one wire has one guard per wire.

Also:
- python/tests/ring_test_distributed.py: main's file. The auto-merge had
  added test_strict_init_after_fallback twice.
- tests/jaccl: the mesh all_gather and sum_scatter calls pass main's new
  stride argument (equal to the chunk size, as main's own
  all_reduce_scatter_gather does).

Tests: tests/jaccl (simulated verbs, ASan and UBSan): 10 of 10 passed.
check_forkdiff.py --upstream-ref upstream/main: 111 changed files, all
described, base afc03c5 is the merge-base. clang-format 21: no changes.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

…ectives branch

Brings in merge/33 at 187b6c6, which is main f78f0fa (the upstream
ml-explore/mlx sync to afc03c5 with its restored history, and the fork
fixes #40, #43-#46) merged into the JACCL send-frame and progress-guard
branch. This branch had merged main up to ad4745a; everything main added
since comes in here.

Conflicts and how they were resolved:
- mlx/fence.h, mlx/backend/{metal,cuda,no_gpu}/fence.cpp, mlx/transforms.cpp:
  main's value-based fence API (upstream ml-explore#4552: update returns a value,
  wait takes it) with this branch's fenced inputs and eager commit.
  Fence::wait_in_consumer now takes the value of the update it waits for,
  and FencedInput carries that value, so a fenced input never waits for a
  later update than its own. The Metal fence keeps the two-word buffer of
  the fused coherent update, now owned through main's allocator::Data.
- Metal fast fence wait: this branch orders the spin with
  wait_for_all_pending (the pass waits for every pass in flight). Upstream
  ml-explore#4689 gets the same order by putting the spin in a pass of its own
  (end_encoding() before it, end_encoding(true) after it). The merge keeps
  the spin in the current pass and adds a buffer barrier after
  wait_for_all_pending, so it also follows the earlier kernels of the same
  pass, as upstream's separate pass does (the same line that the next
  branch, fix/cluster-metal-fence-wait, adds). A wait adds no encoder
  boundary; tensor-parallel decode waits for a CPU-stream all-sum about
  twice a layer. end_encoding(bool wait_in_flight) stays and shares one
  loop with wait_for_all_pending.
- docs/src/usage/environment_variables.rst: both sides' entries
  (MLX_METAL_FUSED_FENCE_UPDATE, MLX_METAL_FENCED_INPUTS, MLX_CONV_WINOGRAD).
- fork.yaml: the merged JACCL branch's file plus this branch's sections
  (startup readiness, two-rank sealed all-sum, fused fence update, fenced
  inputs, cheaper custom-kernel launches). The fenced-input section now
  explains the fence-wait order against upstream ml-explore#4689 and the value-based
  waits. The "not fork work" backport section is gone (upstream ancestry).

Also adapted:
- mlx/distributed/jaccl/lib/jaccl/mesh.h: main's mesh has one MeshImpl per
  wire (upstream ml-explore#4607). The startup wait budget goes to every wire; the
  two-rank exchange of the streamed sealed sum runs on the first wire. The
  two-Mac cluster's device matrix has one device per peer, so it has one
  wire and nothing changes for it.

Unchanged by the merge: main's quantized, matmul and SDPA kernels and their
dispatch (this branch had no change there beyond what main already carried
since #37), main's revert of the upstream ml-explore#4352 simdgroup skip, and this
branch's revert of 4d91320 (main does not carry 4d91320; the footprint
policy keeps {page, page, 0, 0, page}).

Tests (Ultra A, M5 Ultra; Debug build with Metal, 8 jobs):
- tests/tests on the GPU: 310 of 316 cases passed. The 6 failures (linalg
  QR, inverse, Cholesky, pseudo-inverse, eigh, LU) fail the same way on
  main f78f0fa on this machine.
- tests/tests with DEVICE=cpu: 316 of 316 passed.
- ctest process_allocation_cap: 15 of 15 passed.
- check_forkdiff.py --upstream-ref upstream/main: 137 changed files, all
  described, base afc03c5 is the merge-base.

Co-Authored-By: Claude Opus 5.5 <[email protected]>