Refresh MLX fork onto current upstream cutoff - #3

Merged

a-github-name merged 222 commits into

mainfrom

codex/mlx-0.32.1-upstream-refresh

Aug 14, 2026

Commits on Jul 17, 2026

[CUDA][Improvement] RMSNorm forward speed up (#3850)
nastya236
authored
b7c3dd6

Commits on Jul 20, 2026

Fix captured random state in compile (#3828)
angeloskath
authored
ce30733

Commits on Jul 21, 2026

Fix JIT preamble header filter matching project paths containing "Xcode" (#3873)
apocryphx
authored
353440c
Document default value of p in dropout layer docstrings (#3870)
ayaangazali
authored
30a19f7
[WIP] [CUDA] fsdp (#3768)
nastya236
authored
8f64abc
Fix triplet_loss docstring to document the reduced output shape (#3884)
ayaangazali
authored
de82b17
Zero-copy CPU import: mx.array(host_buffer, copy=False) on unified memory (#3872)

HaoXuAI
and
angeloskath
authored
0c537a4

Commits on Jul 22, 2026

Commits on Jul 23, 2026

Commits on Jul 24, 2026

[CUDA] Fix grid overflow in gemm conv unfold kernels for >= 65,536 output positions (#3893)
AdamDLuz
authored
973e27f

Commits on Jul 29, 2026

[CUDA] columnwise quantize with tma (#3157)
nastya236
authored
e6134a8

Commits on Jul 30, 2026

Commits on Jul 31, 2026

Commits on Aug 3, 2026

Commits on Aug 4, 2026

Commits on Aug 5, 2026

docs: Fix missing printoptions page (#3985)
ayaangazali
authored
bd813fb
docs: ThreadLocalStream and iinfo (#3986)
ayaangazali
authored
a681f8b
Make "stop" optional in arange (#3982)
aaishwarymishra
authored
0268c81
docs: Add softsign to the nn functions (#3989)
ayaangazali
authored
5391a8e
docs: Fix broken all_sum and Group references (#3996)
ayaangazali
authored
7c293d1
docs: Fix ast.metal_kernel typo (#3997)
ayaangazali
authored
f99e1dc
chore: Fix unresolved mx.array docstring (#3990)
ayaangazali
authored
391e140
python: fix bfloat16 buffer format itemsize mismatch (#3975)
reckylurker
authored
36e891e
Template Metal complex scalar lanes (#3970)
PhysicistJohn
authored
c4111f8
Pad 2D conv input channels to reach the specialized Metal kernel (#3904)
eyupcanakman
authored
8584610
Support head dimension 96 in Metal full attention (#3943)
dhiltgen
authored
ab946f9
Fix mx.remainder floored-mod for float16/bfloat16 on CPU (#3976)

sashko-zakharchuk
and
zcbenz
authored
2c46b95
Fix sorted gather_mm activation row stride (#3960)
metascroy
authored
268ecd6
Fix state corruption when a primitive throws during eval (#3675)

WindChimeRan
and
zcbenz
authored
fe92a05
python: added __complex__ support (#3984)

aaishwarymishra
and
zcbenz
authored
49ee726

Commits on Aug 6, 2026

Commits on Aug 7, 2026

Commits on Aug 8, 2026

Commits on Aug 9, 2026

Keep randint samples inside [low, high) (#4012)

adityasingh2400
and
zcbenz
authored
5bc4628
chore: Allow saving empty arrays to npy and safetensors (#4080)
ayaangazali
authored
1f9fed2
Fix DeviceType annotations for device and stream arguments (#4059)
XXXXRT666
authored
6a0dd0f
chore: Use dispatch_all_types in DivMod::eval_cpu (#4088)
PhysicistJohn
authored
aadcba1
chore: Use dispatch_all_types in SegmentedMM::eval_cpu (#4087)
PhysicistJohn
authored
62e2508
chore: Use dispatch_all_types in Arange::eval_cpu (#4085)
PhysicistJohn
authored
0c53491
Fix Device and Stream lexicographic ordering (#4086)
wgu9
authored
e2aa0d0
Enable complex64 scatter addition on GPU (#4078)
PhysicistJohn
authored
b0c3562
Optimize large NVFP4 QMV on M5 Max (#3961)

dhiltgen
and
zcbenz
authored
5a1e44c
Fix ring hanging on peer disconnect (#4060)

erwinzhang7
and
zcbenz
authored
8c28c38
Propagate NaN in cummax and curmin (#4047)

JasonHonKL
and
zcbenz
authored
7eb3d85

Commits on Aug 10, 2026

Commits on Aug 11, 2026

Commits on Aug 12, 2026

docs: Writing a Fast KV Cache (#4019)

authored
2ddcb6f
Fix and enable non-transposed NAX qmm (#4051)

gordofreemo
and
zcbenz
authored
d04e5db
Add scatter reduce for JACCL (#3901)
angeloskath
authored
31b5cbb
Fix diag for zero-size input (#4165)
devteamaegis
authored
56c26e8
Add an inverse-CDF path to categorical sampling (#4177)

deescondrillas
and
zcbenz
authored
13e7968
Fail when a requested RDMA device is not found (#4180)
erwinzhang7
authored
307692c
Bound GGUF tensor data offsets against the file mapping (#4179)
robertomeroni
authored
ae8e9a7
Shift by the max in log_softmax (#4169)
ayaangazali
authored
d935276
Shift by the max in cross_entropy (#4188)
ayaangazali
authored
74c9acd
Fix linalg.norm(x, keepdims=True) not keeping dims (#4166)
devteamaegis
authored
152baf4
chore: Use dispatch_inexact_types in binary_float_op_cpu (#4194)
PhysicistJohn
authored
a435674
Implement batched matmul for large 1D dot products (#3580)
Ved235
authored
36fde27
Re-disable qmm_n_nax and fix group_size < 64 (#4202)

PhilipJohnBasile
and
zcbenz
authored
01d4e12
chore: shell-quote the working directory in the distributed launch script (#4189)
bunlongheng
authored
d34f630
Report when no usable GID is found (#4191)
erwinzhang7
authored
21d897d
Keep double precision for python floats in float64 operations (#4173)
nileshpatil6
authored
3abd0fd
Jaccl ring multi wire fix (#4193)
angeloskath
authored
e3d389e
Report why creating a queue pair failed (#4209)
erwinzhang7
authored
52960f8

Commits on Aug 13, 2026

Commits on Aug 14, 2026

Add endpoint parameter to linspace (#4184)

AKnassa
and
zcbenz
authored
1d717bd
Fix vmap of partition/argpartition dropping the kth argument (#4116)
Adityaj0
authored
306bdcd
Fix nan_to_num replacing inf with 0 for float16 and bfloat16 (#4222)

codeAnqiang-ma
and
zcbenz
authored
d9ad465
Fix einsum not broadcasting batch dimensions in batched tensordot (#4125)

Adityaj0
and
zcbenz
authored
bd5c3a2
Add 1-bit affine quantization support

khosravipasha
authored and
a-github-name
committed
d2d7138
Guard fast-path Metal kernel dispatch for 1-bit quantization

khosravipasha
authored and
a-github-name
committed
5743cb8
metal: 1-bit qmv_fast use 1 pack/thread for occupancy (#3)

bri-prism
authored and
a-github-name
committed
3e0fa1f
Disable NAX on generation 17 for low-bit correctness
a-github-name
committed
564c116
Route affine qmv_wide by bit-width and batch size

bri-prism
authored and
a-github-name
committed
629f091
Add native affine 1-bit CUDA QMV
a-github-name
committed
3ebf620
Format rebased low-bit changes
a-github-name
committed
dde9c83
Route low-bit CUDA matmuls through native QMV
a-github-name
committed
f8afc8b
Keep global CUDA streams alive through process exit
a-github-name
committed
d285f72
Expose quantized headers to custom Metal kernels
a-github-name
committed
9190776
Speed up NVFP4 block-loader staging
a-github-name
committed
4d88c9d
Tune large H3 GEMMs on M4 Max
a-github-name
committed
0dfe6cb
Align 1-bit qmv fast dispatch geometry
a-github-name
committed
3a0382c
Format refreshed MLX patches
a-github-name
committed
b34c235
Link quantized custom-kernel sources in AOT builds
a-github-name
committed
2ccca66
Make generated NAX optimizations source-reproducible
a-github-name
committed
1fb504c
Fix affine and fallback QMV template dispatch
a-github-name
committed
3f6e772
Fix Swift 6.0 Linux matmul compilation
a-github-name
committed
b57bd76