Break the sibling cycle when an array is released by assignment by tudalex · Pull Request #4453 · ml-explore/mlx

tudalex · 2026-09-02T17:50:19Z

This PR fixes a memory leak issue that I encountered on iOS when unloading Qwen models.
With the help of AI I tracked it down to the compile function and further down to a corner case in reference counting in MLX's Array class.

When reassigning an array that is produced from a multi output (e.g. the result of a split call), the destructor code is not called, so there is no check to see if it is the last reference to a family.
This fix just makes sure that the destructor is actually called on the previous value when reassigning an array, thus fixing the memory leak.

I think this should also fix issue #3932.

A 4B params Qwen model would leak ~1Gb each time it is unloaded which is quite a lot for a mobile device.

☑️ I understand it is strictly prohibited to use AI to write PR description
AI usage disclosure:
Used to track down the bug and write tests that could reproduce it. All the tests are generated by AI and fail on the main branch. Commit description also generated by AI. Multiple solutions also brainstormed with AI, I considered this one to be the cleanest potential fix.

zcbenz

Thanks for the fix and I think you found the root cause, but I need to think for a while about whether there is a better fix that does not requiring patching every place copying array_desc.

About the tests though, using get_active_memory and gc.collect() in tests is very unreliable and was causing a lot of flakes. Can you turn them to C++ tests like this one

mlx/tests/array_tests.cpp

Lines 675 to 692 in 90846ad

    
           // https://github.com/ml-explore/mlx/pull/1590 
        
           TEST_CASE("test siblings circular references without eval") { 
        
             std::weak_ptr<array::Data> tracker; 
        
             auto fun = [&]() { 
        
               array key({1, 2}); 
        
               auto splits = split(key, 2); 
        
               { 
        
                 // Set fake data as a tracker for ArrayDesc's lifetime. 
        
                 splits[0].set_data(allocator::malloc(0)); 
        
                 tracker = splits[0].data_shared_ptr(); 
        
               } 
        
               auto a = reshape(splits[0], {}); 
        
               auto b = reshape(splits[1], {}); 
        
               return b; 
        
             }; 
        
             fun(); 
        
             CHECK(tracker.expired()); 
        
           }

zcbenz

Thanks for updating the tests, I think the PR is ready to go with just a little refactoring.

tudalex · 2026-09-03T06:36:52Z

Thanks for the fix and I think you found the root cause, but I need to think for a while about whether there is a better fix that does not requiring patching every place copying array_desc.

Like having the array_desc_ destructor trigger this? I'm trying to see if I can make that work, would be even cleaner and prevent future bugs.

tudalex · 2026-09-03T07:31:33Z

Here is an alternative in which we wrap the array_desc_ in a class, and it's destructor triggers the reset logic tudalex@3602800
let me know if you like this one more. To me it seems cleaner from an OOP viewpoint, but it also adds one more level of indirection for anybody reading the code.

zcbenz · 2026-09-03T09:41:56Z

Thanks for experimenting with the ideas and sorry I'm being nitpicking here, I would like to keep the code self-explanatory so anyone even without too much C++ expertise can still understand and maintain the code, that's why I don't like the original array previous(std::move(*this)) trick or the DescRef abstraction.

What do you think about zcbenz@a249d99? I think as long as we are not doing a self-assignment there shouldn't be need to keep an additional copy of assignee?

zcbenz

Many thanks for working out the edge cases!

The outputs of a multi-output primitive hold each other in `siblings`, a reference cycle that `~array()` breaks once it sees the last external reference go. Assignment did not run that check: it replaced the descriptor directly, so an array whose last external reference went away by assignment left the cycle intact, and with it every input below it in the graph. `compile_dfs` does exactly that. It rewires the outputs' inputs to the copied tape by assignment, which drops the trace's own consumers of a Split (or any multi-output primitive), and then rewrites the parents map by assignment, which drops the last reference to the trace's outputs of that Split. From then on the cycle owns the whole trace beneath it, captured constants included, and nothing ever frees it: `compile_erase` and `compile_clear_cache` never reach it. Compiling a function that splits a matmul of a captured weight leaked that weight once per compiled function. In mlx-swift-lm, whose Qwen3.5 decode step is compiled and splits the fused q/k/v projection in every layer, that was 0.5-1.7 GB of a 4B model left resident after the model was dropped, which killed an app that loaded another model next. Route copy assignment, move assignment and `overwrite_descriptor` through the destructor by moving the previous value into a temporary, so the cycle check runs on every way an array reference can end. `a[...] = v` on the last surviving output of a lazy split reaches the same bug with no compile involved: it goes through `overwrite_descriptor`, and with the other output already gone the two held only each other, and the split's input with them. There is a test for each path.

get_active_memory() and gc.collect() are unreliable in a test and have caused flakes, so track the ArrayDesc's lifetime with a weak_ptr to its data instead, in the style of the existing "test siblings circular references without eval" case. One test drops both outputs of a lazy split by assignment, covering the copy and the move operator; the other leaves a single output alive and replaces it through overwrite_descriptor, which is the path __setitem__ takes and where no compile is involved at all. Both fail on main and pass with the fix, and neither measures memory or touches the garbage collector.

The destructor body becomes array::reset(desc): it swaps desc into array_desc_ and runs the check on the descriptor it lets go of. The destructor, both assignment operators and overwrite_descriptor call it instead of releasing through a temporary array. reset() holds the new descriptor before checking the old one. The order matters when a sibling is assigned over the last outside reference to its own cycle, because the sibling is only reachable through that reference: checked first, the pair looks unreferenced, the cycle is broken and the sibling list the new value lives in is freed before the assignment reads it. Held first, the new descriptor counts as an outside reference and the cycle stays intact, as it did when the check ran after the assignment. The added test segfaults with the other order.

reset(desc) becomes a static release(desc): it takes the descriptor the array is letting go of and runs the sibling-cycle check on it, touching nothing else. Callers hand it the old descriptor with std::exchange, which is also what takes the new one first: release(std::exchange(array_desc_, other.array_desc_)); The order is the fix, not the naming. A reset() reading array_desc_ cannot provide it: the member still holds the old descriptor while the check runs, so nothing holds the new one, and when the new one is a sibling of the old the pair looks unreferenced. The cycle is then broken and the sibling list the new value lives in is cleared out from under it, leaving an output with no siblings for outputs() to index and for Split::eval to write. Comparing array_desc_ instead of id() does not catch this, since the two descriptors genuinely differ: they are different outputs of the same split. Self-assignment is a separate case and both guards already cover it. std::exchange copies the new descriptor once and moves the old one out, so this is the single atomic increment a plain copy-assign already does. The assignment operators also use the array_desc_ comparison now, which is cheaper and clearer than comparing id().

…xplore#4453) Co-authored-by: Cheng <git@zcbenz.com>

The assignment case is ml-explore/mlx#4453, merged 2026-09-11 and not in the pinned mlx; the compiled split case is ml-explore/mlx#3932, open.

…ling probes) The assignment-over-siblings leak is ml-explore/mlx#4453, merged upstream on 2026-09-11 and not yet in the pinned mlx; the compiled split leak is the open ml-explore/mlx#3932.

MLX 0.32.2's mx.compile leaks every multi-output primitive that is an intermediate of a traced graph: compile_dfs rewires the trace by assigning over the old outputs' handles, and array assignment skips the sibling-cycle break that ~array runs (fixed upstream by ml-explore/mlx#4453, not in 0.32.2), so the group's descriptors keep each other alive together with every array they reference. The one-token FFN half that Glm5NextDecoderLayer compiles runs through multi-output fused kernels whose inputs are the layer's weights: since 4a9d32e the route-selecting gate/up kernel (act, indices, scores) reads the routed and shared gate/up weights, and the deferred HC chain's pre/post kernels read the connection weights. As trace constants those weights stayed allocated after the model was dropped: each compiled MoE layer could leak its ~2.5 GB of gate/up experts, so unloading GLM-5.3-Flash-oQ4e after a request left 43-77 GB active (dev server: "Settle barrier timed out ... freed=108.53GB", "Emergency reclaim failed ... active_memory=63.85GB"), and the next load was killed for memory. Upstream main's compiled FFN leaks only the small HC mixing weights this way (up to ~1 MB in the test model). compile_ffn_block() compiles the FFN methods with the arrays of the modules they read (ffn_hc, post_attention_layernorm, mlp) passed as inputs: in the trace they are shape-only placeholders, so a leaked group holds no weights. The same arrays are passed on every call, so the graph, the kernels and the values are unchanged (about 0.5 us more per compiled call for the extra inputs). The MTP verify runtime compiles its FFN through it too. GLM-5.3-Flash-oQ4e on the combined dev tree (e2105bf3 + this commit), engine path in one process, engine_pool-style unload, active memory left: admin-bench requests (pp 32/1024/4096/16384/65536, tg 128): 56.19 GB -> 0.00 GB 1024-token prompt + 2x255 decode tokens + 16k prompt: 58.74 GB -> 1352 B stock mlx 0.32.2 wheel, 1024-token prompt + decode: 61.30 GB -> 1336 B main + perf/glm-decode-2 alone (1024 + 32 tokens): 56.19 GB (main: 0.01 GB) Greedy decode at 1024 tokens 61.2/61.3 -> 61.7/61.8 tok/s, 16k 55.7 -> 56.0 tok/s, identical output. Small GLM-5.3 model of the decode-kernel tests (prefill + 3 one-token steps, 3 runs): 73-220 MB left -> 0-36 bytes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The speech and language models now load when Casper launches and stay loaded. A model unloads only when the user switches to another one, or, for the language model, when refining is turned off; a model in use by a dictation unloads once that dictation finishes. The Keep Model Alive setting is gone. Benchmarks showed loading while the user speaks hides the load for recordings longer than about 0.8 s, so the timed options saved little, and each language model reload leaks memory in mlx-swift (fixed upstream in ml-explore/mlx#4453, not yet released). Keep Microphone On is unchanged.

…d HC chain, routed gate/up) (#4019) * perf(glm5_next): fuse decode/verify op chains into exact kernels Single-token decode and short verify forwards (B=1, L<=8) of GLM-5.3 ran ~6600 dispatches per token: MLX encode on the host (~3.8 us/dispatch) and the GPU were both ~22 ms per step. Replace the longest op chains with JIT kernels that reproduce the reference arithmetic bit for bit: - HC: hc_mix (fp32 RMS + mix GEMV, gemv reduction order) + exact_hc_norm replace cast/RMSNorm/GEMV/sinkhorn/RMSNorm; one-token hc_expand issues the same NAX relaxed-precision matmul2d as MLX's fp32 GEMM (only when MLX itself would: NAX available and MLX_ENABLE_TF32 on) plus the epilogue. - MoE: one-token router (gemv + sigmoid/bias, stable-sort top-k), fused routed+shared gate/up SwiGLU and down+weighted-combine kernels for route counts SwitchGLU leaves unsorted (L*topk < 64). - KDA: one kernel per layer between in_proj and o_proj (conv, SiLU, l2norm, qmv_quad gate projections, safe gate, delta rule, RMSNormGated). - DSA indexer: decode score kernel equal to the Steel tile + mask, and the top-k index expansion (fixes an unusable .strides access). Exactness: 42 unit tests compare every path bitwise with the reference op graph, including an end-to-end small model (decode and verify widths 1..8, with and without TF32 NAX). Engine greedy decode, 256 tokens at pp200/1024/4096: token-identical to OMLX_GLM5_DECODE_FUSION=0. Measured (M5 Ultra, engine greedy, same tree, fusion off -> on): 38.9/38.7/32.7 -> 51.7/51.2/49.0 tok/s at pp200/1024/4096. Raw forward at ctx 200/4096: L=1 24.8/29.8 -> 18.9/20.3 ms, L=4 39.2/45.8 -> 32.3/35.7 ms, L=8 50.3/59.3 -> 42.5/48.8 ms. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): exact verify-block router and wider small decode kernels With the op chains fused, one-token decode is GPU-bound (the host waits on in-flight work ~half the time) and the remaining GPU cost is dominated by weight streaming plus a few latency-bound small kernels. Make those use more of the GPU and fuse the verify router, without changing a bit: - hc_mix: one mix row per 256-thread threadgroup (24 threadgroups per token instead of 6), replaying the 1024-thread rms_looped reduction tree. - one-token router logits: one expert row per simdgroup (the gemv's per-row arithmetic does not depend on rows per thread). - router select: rank-based top-k (one thread per expert, stable-sort tie order, NaNs last) instead of eight serial selection rounds. - verify blocks (2..8 tokens): the router logits reproduce MLX's NAX split-K fp32 GEMM (same relaxed matmul2d chain per K partition, same partition sum), followed by the fused sigmoid/bias/top-k/normalize. Used only where MLX itself runs fp32 GEMMs on NAX with TF32 on. 44 bitwise tests (TF32 NAX paths checked in a TF32-enabled subprocess, including an end-to-end small model with the verify router engaged). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): fold the verify-block shared expert into the MoE gate/up kernel For 2..8-token verify blocks the shared expert's gate and up projections ran as two MLX qmv_wide matmuls plus five clamp/SiLU/multiply kernels per MoE layer. The fused gate/up kernel gets one extra slot that replays qmv_wide's multi-row arithmetic (8 lanes per row, per-group dequantize in 8-value sub-chunks, 4/2/1 shuffle ladder) for all tokens at once and applies the clamped SwiGLU, so the shared weights are read once per layer and seven dispatches become none. The shared down projection keeps MLX's own qmv_wide call. Bit-identical (MoE tests at L=2/4/7 with 8- and 4-bit shared experts assert the new slot is used). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): exact fused latent attention for DSA decode and dense verify GLM-5.3's NoPE MLA attends in the 512-wide latent space, which mx.fast.scaled_dot_product_attention serves with its op fallback: a bf16 scale multiply, a NAX GEMM for the scores with the 64 heads folded into M, a where() mask, a precise softmax, a second NAX GEMM (split along K past ~1k keys) and, for one-token sparse decode, a materialized gather of the 2051 selected latent rows. The small-M GEMMs run on very few threadgroups. Three kernels replay that exactly: the same bf16 16x32x16 relaxed matmul2d chains per K partition (steel_matmul's split-K routing and partition sums, two 16-wide steps per 32-wide k step, zero-padded tails), the softmax_single_row reduction tree, and the same rounding points. Sparse decode reads the selected (clamped) rows in place. Used for one-token decode (dense and sparse) and dense 2..8-token verify blocks with a bool or causal mask when NAX is available; sparse verify keeps _gathered_attention. Tests: 46-case bitwise sweep against SDPA (1..3000 keys, both split-K regimes, sparse gathers with negative indices, L=2..8 masks), plus the end-to-end small models (dense and sparse prompts, TF32 on and off). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): exact sparse latent attention for DSA verify blocks Past 2048 keys a 2..8-token verify block runs _gathered_attention: each token gathers its 2051 selected latent rows ([L, 2051, 512] copies) and the SDPA fallback solves L * 64 one-row problems with MLX's batched gemv (scores) and gemv_t (values), re-reading every token's rows once per head. Replay those reductions exactly (gemv: 4 keys per simdgroup, 16 products per lane, 16/8/4/2/1 ladder; gemv_t: 8 key groups x 4 lanes, 4 simdgroups per 64 columns, leftover keys, 16/8/4 ladder; the shared precise softmax) while reading each selected row in place, once per chunk of 8 heads. Tests: 12 bitwise cases against the gathered SDPA (L = 2..8, 64..2051 selected keys, negative/clamped indices) and the sparse end-to-end small model now exercising the path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): one dispatch per input for the DSA attention projections Decode/verify blocks (B == 1, L <= 8) of the DSA layers now compute the projections that share an input in one multi_qmv dispatch per quantization: q_a, kv_a, the indexer's wk (and weights_proj when the indexer selects) from x, and q_b (with the indexer's wq_b) from qr. multi_qmv replicates qmv_fast for one token and qmv_wide for 2..8 rows per output row, so every projection is bitwise the separate matmul; the indexer takes the precomputed projections. Layers whose projections do not share (bits, group size) keep the reference calls. The one-token router selection goes back to the simdgroup top-k rounds (11.5 us vs 16.2 us for the per-expert rank kernel on GLM-5.3) and now orders NaN scores after every number, lowest index first, like argpartition's sort. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * chore(glm5_next): per-family switch for the fused decode kernels OMLX_GLM5_DECODE_DISABLE (comma-separated STATS keys) or decode_kernels.DISABLED at runtime sends the named families to the reference path, so each kernel's contribution can be measured in one process on the same weights. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): prefetch the latent attention NAX operands The latent scores/values kernels waited on every operand load before each 16x32x16 NAX op (and, for sparse decode, on an index load before the value row load), so a 2048-key partition took ~130 dependent memory round trips. The operands of the next 8 chunks are now loaded ahead of the accumulator chain (the chain and every op are unchanged), selected value rows are staged in threadgroup memory, and the key rows are addressed by offsets instead of an array of device pointers. The Metal compiler produced wrong code for a 4-deep ring in some split-K configurations, so each compile-time configuration is checked bitwise against the SDPA fallback on its first call and falls back to the reference if it differs; tests cover every configuration GLM-5.3 decode and verify use, and several ring depths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): prefetch the verify router's NAX operands Same treatment as the latent kernels: the split-K router logits chain (128 dependent 16x32x16 TF32 ops per partition) now has the operands of the next 8 chunks in flight, and rows are addressed by offsets. Each configuration is checked against group_expert_select on its first call and falls back to the reference if it differs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): route-slot-major MoE grids for verify blocks Verify blocks route 2..8 tokens to overlapping experts (about 24 distinct of 32 routes at 4 tokens, 40 of 64 at 8), but the gate/up grid ran one route slot after another (and the down grid one token after another), so an expert shared by two tokens was streamed from DRAM twice, hundreds of microseconds apart. For blocks of 2+ tokens the slots (tokens) now vary fastest in the grid, so shared experts' row blocks are read back to back. Every threadgroup computes exactly what it did before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): fold the verify-block shared expert down into the combine For 2..8-token blocks the shared expert's down projection ran as its own [T, 2048] x [4096, 2048] quantized matmul before the routed combine kernel added it. The combine kernel now computes it in place for each (token, row) with the multi-row qmv_wide arithmetic that matmul uses (8 lanes per row, per-token accumulators independent of the other tokens), then adds it exactly as before: one dispatch and one [T, 4096] round trip less per MoE layer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): keep latent attention off MLX's gemv-routed score shapes The fused latent attention replicates the SDPA fallback's NAX GEMMs, but MLX computes the scores q @ k.T with gemv_wide when the folded query has 2..15 rows (heads x tokens) and with gemv for a single key, so with a few-head model (the 8-head test model at one token) a handful of outputs differed by one bf16 ULP on some inputs (seed 48 / 300-token prompt, and the combined build's KDA-prefill activations). GLM-5.3 (64 heads, >= 64 rows) never takes those routes; the kernel now declines them. Tests: the small model uses 16 heads (so the one-token sparse path stays covered), a seed x prompt-length x heads sweep (including the failing seed 48 / 300 / 8 heads case), explicit declines, and a KDA-fast-prefill then fused-decode check that runs when that prefill is in the build. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): fused decode experts also read a fused gate_up_proj With the MoE gate/up fusion (perf/moe-fused-gate-up) GLM-5.3's SwitchGLU holds one gate_up_proj ([E, 2N, K]: gate rows then up rows per expert) instead of gate_proj/up_proj, and _decode_experts failed on the missing attributes. The gate/up kernel now takes the expert row stride and the up-row offset as template constants, so the fused tensor is read in place (passed as both operands, up rows N further on) with the same per-row arithmetic; the sort threshold uses the projections SwitchGLU itself passes. Tests: fused layout vs split layout vs reference, bitwise, and vs the fused SwitchGLU forward where the build has it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): leave the fused latent attention off by default In the model the latent attention replicas (one-token / dense verify NAX kernels and the sparse verify rows kernel) cost more than the reference ops they replace: with one call in flight each output tile is a chain of up to 128 dependent 16x32x16 NAX ops at low occupancy (single-call latency e.g. ~300 us vs ~60 us for the 2051-key one-token reference, ~830 us vs ~260 us for a 4-token sparse verify; disabling them made the in-model one-token step 1.5-2.9 ms faster at 1k-4k context and the 4/8-token verify 3.7/6.7 ms faster at 4k). They stay exact and tested, and OMLX_GLM5_DECODE_DISABLE (which now replaces the default-off set) re-enables them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): skip the first-use kernel checks inside mx.compile The latent attention and verify-router first-use checks evaluate a comparison with .item(), which raises while tracing (mx.compile / vmap). When that happens the call now takes the reference path without caching a verdict, and the check runs on the next eager call. (The verify router only runs eagerly in the model; the compiled FFN covers one-token blocks, whose router has no check.) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(glm5_next): fused decode after upstream's fused KDA prefill The small model's 512-token prompt chunks now run #3944's fused KDA prefill prework and norm-gate (glm53_kda_prework); the fused decode and verify paths stay bitwise equal to the reference decode on the caches it writes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(glm5_next): router tests pass on the stock mlx 0.32.2 wheel Two router tests only passed on a source build of mlx. On the stock 0.32.2 wheel (what oMLX ships): - the one-token router stays bitwise at GLM-5.3's shape (288 experts x 4096), but for the small 64 x 512 test shape the wheel's precompiled gemv uses another reduction variant than the one the kernel replays; that shape now allows the summation-order difference (<= 2 fp32 ulp, identical expert indices); - with TF32 NAX matmuls on, the verify-block router's first-use check declines one of the five row counts (its TF32 GEMM rounds differently there) and that configuration keeps the reference path, as designed. The test now accepts a declined configuration; every call is still checked bitwise against the reference. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(glm5_next): skip model-level fused-decode checks under the MTP runtime glm5_next_vlm_runtime.apply() (MTP checkpoints) is process-wide and replaces the layer __call__ methods, which do not route the fused decode kernels. Tests that ran it earlier in the process (e.g. the oQ round trip with nextn weights) made seven model-level checks here fail when they ran first; those checks now skip in that case. Kernel-level tests are unaffected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): fused decode kernels default to M5 (NAX) GPUs The fused decode/verify kernels replay the reference op graph bit for bit on M5 GPUs, where they were validated. On other GPUs the reference ops run different MLX kernels, so the replay is not bitwise, and several kernels use 1024-thread threadgroups that some devices do not allow (the GitHub macOS runners' virtual GPU caps them at 768 or 640 and failed the CI). OMLX_GLM5_DECODE_FUSION now defaults to on only when NAX is available (=1 forces the kernels on, =0 turns them off everywhere); the bitwise kernel tests skip on non-NAX devices. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): KDA decode sigmoids match the eager kernels on release wheels The fused KDA decode kernel evaluates beta = sigmoid(b) and the RMSNormGated gate sigmoid with MLX's Sigmoid expression, whose metal::exp is compiled with the runtime default math. The reference computes both with eager mx.sigmoid calls, and the release mlx wheels run those from the precompiled metallib, which is built with -fno-fast-math (precise exp). On the stock mlx 0.32.2 wheel the two differ in ~20% of sigmoid values; after bf16 rounding the kernel's outputs differed from the reference in a few elements for 6 of 10 test seeds (then in the recurrent state from the next verify block on). Source builds that JIT every kernel use the default exp for mx.sigmoid too, which is why the existing checks passed there. eager_sigmoid_precise() now decides once per dtype which expression the eager kernel reproduces (sweep of 2^17 inputs; None when neither, and the fused path then declines), and the kernel selects exp for those two sites accordingly. Every other exp in the fused kernels has a compiled-graph or custom-kernel reference (nn.silu, compute_g_safe, group_expert_select, the HC kernels) and keeps the runtime-compiled expression. New tests: the probe reproduces mx.sigmoid, and a 10-seed decode/verify sweep of the KDA layer is bitwise equal to the reference (6 seeds failed on the stock wheel before this change). 168 passed on the stock wheel and on the source build. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 98cc9ba245ec6fe4cae7fceed87e33ec45bff33f) * fix(glm5_next): one-token router sigmoid matches the eager kernel on release wheels group_expert_select takes the sigmoid of the fp32 router logits with MLX's eager Sigmoid kernel (it is not fused into a compiled kernel there). On the release wheels that kernel is precompiled with -fno-fast-math and uses the precise exp, while the fused router kernels used the runtime compiled exp: a last-bit difference for a few percent of logits that, after the top-k normalization, changed ~5% of the routing weights by one ulp. The decode then left the reference path at the first decode step (stock wheel, 4096-token context: max logit difference 0.44, tokens diverging at step 11; source builds were exact). This predates perf/glm-decode-2 (perf/glm-decode's one-token and verify routers). The router logits kernels now take the precise sigmoid where eager_sigmoid_precise(float32) says the eager kernel does (probed once, eagerly, from the decoder layer before its compiled FFN block is traced; undecided means the reference router). Test: the one-token router against group_expert_select over 96 draws per shape, including the checkpoint's router bias level (the old test tolerated 2 ulp on one shape and passed by chance with 12 draws on the other). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 204d8277950adb3d2546693dca6329cc31b5fc34) * perf(glm5_next): start evaluating one-token decode steps early A GLM-5.3 decode step is ~800 dependent dispatches and its Python graph build takes ~2.7 ms. mlx allows ~10 command buffers in flight, and with its default per-buffer size budget (every expert or projection weight input counts in full) one step commits ~170 buffers, so while the next step's graph is being built the GPU drains the few buffers still queued and then idles: ~1.3 ms per token at 1k context, ~1.6 ms at 4k (devbuild timeline: "host-starved gaps"). One-token forwards now call mx.async_eval on the hidden state every 8 layers (OMLX_GLM5_DECODE_EVAL_EVERY, 0 = off), so the first layers are encoded and running while the later ones are still being built. The computed values are unchanged (scheduling only); the MTP runtime's replacement forward follows the same policy. Raw decode steps, GLM-5.3-Flash-oQ4e, M5 Ultra, mlx's default buffer budget: 1k context 17.99 -> 16.83 ms/token (55.6 -> 59.4 tok/s), 4k 19.45 -> 17.83 ms (51.4 -> 56.1). With an op-count-only buffer budget (no host gap to hide) it is neutral: 16.72 -> 16.67, 17.77 -> 17.71. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): fused KDA decode for layers with mixed projection bits The fused KDA decode step needs the q|k|v|f_a|g_a|b projection of the layer as one row block. It was only built when all six projections share one quantization, so a layer whose v_proj is 5-bit among 8-bit q/k/gates (GLM-5.3-Flash-oQ4e layer 40) decoded through the reference body: six projection matmuls, concatenations, the short conv, SiLU, l2norms, gate projections, the gated-delta kernel and RMSNormGated as ~35 dependent dispatches instead of three. Such layers now get one fused weight per quantization (a single-member group uses the layer's own arrays) and the decode projection runs one matmul per group, putting the rows back in q..b order with one concatenation of contiguous runs. Each row is the row of its own layer's matmul (same qmv kernel and per-row arithmetic), so the block equals the reference projections and the fused kernel sees the same input. Test: a KDA layer with a 5-bit or 4-bit v_proj decodes and verifies bitwise like the reference and takes the fused kernel. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): 5-bit KDA gate projections inside the decode kernel The fused KDA decode kernel replays the forget/output gate projections (f_b, g_b: K = 128 rows) for 4- and 8-bit weights, which MLX runs as qmv_quad. GLM-5.3-Flash-oQ4e quantizes them to 5 bits in 8 of its 34 KDA layers; there the two projections ran as separate quantized matmuls before the kernel: MLX's one-row qmv with only 16 of 32 lanes active for K = 128, two latency-bound dependent dispatches (~20 us each in a one-dispatch-per- buffer timeline) per layer and token. For one token the kernel now computes these rows itself with qmv's arithmetic: lanes 0..15 load 8 values each (load_vector_safe/qdot_safe with N = 8, identical to load_vector/qdot), the other lanes contribute zero, one simd_sum per row. Verify blocks (2+ rows go to qmv_wide) keep the precomputed projections. Tests: the 10-seed decode/verify sweep now also runs with 5-bit gates and checks that one-token steps take the in-kernel rows; 322 passed on the stock wheel and on the source build. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): select one-token MoE routes inside the gate/up kernel A one-token MoE block ran the router as two dependent dispatches (logits GEMV, then a single-simdgroup top-k select) before the gate/up kernel. The routed gate/up threadgroups now replay the select kernel's loop on the biased sigmoid scores themselves (the same code, moved into a shared helper: stable descending order, ties to the lower expert index, NaNs last), and slot 0 / tile 0 publishes the routes and normalized routing weights for the down/combine kernel. The select dispatch disappears; the selection is ~1 us of simdgroup work at the start of each routed threadgroup. Verify blocks keep moe_router_rows + select. OMLX_GLM5_DECODE_DISABLE= router_select_fused restores the two-kernel router. Test: routes and outputs bitwise equal to the router select path and to the reference MoE for GLM-5.3's shape (288 experts, K 4096), including exact score ties. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): overlap the one-token router with the shared expert In a one-token MoE block the shared expert's gate/up does not depend on the router, but it ran inside the routed gate/up kernel, after the router logits kernel: a 4.7 MB fp32 GEMV spread over 72 threadgroups that is mostly latency. The shared expert's gate/up now runs as its own dispatch (the shared slot of the fused kernel, same code and arithmetic) ahead of the routed one. Neither it nor the router logits kernel reads the other's output, so mlx encodes no barrier between them and they run concurrently: the router's latency hides under the shared expert's 17.8 MB weight stream. The routed gate/up (with the in-kernel top-k) takes the shared activation as an ordering-only input, and the down/combine kernel reads it from its own buffer (moe_down_combine(shared_act_sep=...), same qmv arithmetic as the in-act slot). OMLX_GLM5_DECODE_DISABLE=moe_shared_split restores the single gate/up launch. Test: bitwise equal to the single-launch path, the router-select path and the reference MoE for GLM-5.3's shape, including exact score ties. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(glm5_next): model-level checks require the one-token router paths The small-model decode/verify logits checks (bitwise against the reference model, dense and sparse DSA attention, seed sweep) now also assert that the in-kernel route selection and the separate shared-expert gate/up launch were taken. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): one-token dense MLP gate/up and SwiGLU in one dispatch GLM-5.3's first three layers use a dense MLP (8-bit, 4096 -> 12288). For one token its gate and up projections ran as two quantized matmuls (concurrent, ~0.89 TB/s together for 100 MB in a one-dispatch-per-buffer timeline) followed by the clamped SwiGLU as a third dispatch. They now run as one dispatch of the shared-expert gate/up kernel (Glm5NextMLP is the shared expert's module; same one-row qmv arithmetic per row and the same clamped-SwiGLU epilogue), then the reference down projection. Test: bitwise equal to the eager and the compiled reference MLP (8- and 4-bit). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): bitonic top-k for decode/verify indexer rows Past index_topk tokens of context the DSA indexer picks 512 of up to 2048 pooled blocks per decode/verify row with the native radix-select kernel: two histogram passes whose threadgroup atomics contend on the few bins that hold similar scores, and two 256-bin scans by a single thread (~15-20 us per sparse layer in a one-dispatch-per-buffer timeline). Its output is a function of the scores alone: keys are the 16-bit order-preserving score bits, strictly greater keys fill the output in index order, then threshold ties in index order. dsa_topk_rows finds the threshold with a threadgroup bitonic sort of (key, index) pairs (one 1024-thread threadgroup per row, P <= 2048) and emits the same output with the native kernel's raking index-order pass. Rows beyond 2048 blocks and other select widths keep the native kernel. Tests: identical indices and order to the native kernel for 512..2048 blocks, 1/4/8 rows, exact ties, -1e30 masked blocks and a NaN; the small-model sparse decode/verify logits check requires this path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): load indexer query fragments ahead of the score MMAs dsa_decode_scores loaded each 8-wide query fragment inside the per-head chain of 16 accumulator-dependent simdgroup MMAs, so every step waited on its own load. Each head's fragments are now loaded before its chain; the MMA sequence (and so every score) is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): fold each one-token HC expand into the next HC pre A one-token half-layer ran three dependent, latency-bound HC dispatches around its branch: hc_expand of the previous branch output, hc_mix (RMS + 24-row mix GEMV) and exact_hc_norm (sinkhorn, collapse, RMSNorm). The branch input only needs the four pre rows of the mix, so: - hc_pre_fused computes the mix rows four per 1024-thread threadgroup (hc_mix's multi-row layout) and threadgroup 0, which owns the pre rows, finishes exact_hc_norm's collapse and RMSNorm itself: one dispatch on the chain instead of two; - hc_post_mm (post weights, sinkhorn comb and hc_expand_one's NAX comb product of the residual) only feeds the next expand; the branch input depends on it without binding it, so it is scheduled first and runs beside the branch's first kernel; - the next hc_pre_fused applies hc_expand_one's epilogue to the branch output while loading h and stores h (the next residual), so the expand dispatch disappears. The model loops (vendor and MTP runtime) pass the pending expand between layers and materialize the last one. Every value keeps the reference kernels' arithmetic and order (the sinkhorn code is compiled ahead of the MPP include, as in exact_hc_norm). Tests: a chained kernel check against HyperConnection + RMSNorm + hc_expand (hidden 4096/2048/1024, TF32 NAX on) and the small-model decode/verify logits, also under the MTP runtime loop, bitwise. Measured (stock wheel): in a dependent chain of 90 HC pres alone, 26.0 -> 20.6 us per half-layer; with a two-qmv branch per half-layer (71 MB of 8-bit weights) 103.6 -> 100.0 us. Raw pipelined decode of the model with this commit and the next two (MLA head qmv, sparse row gather): 16.81 -> 16.72 ms/step at 1k context, 17.94 -> 17.62 ms/step at 4k. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): one-token MLA embed_q/unembed_out in wider threadgroups QuantizedMultiLinear runs the per-head MLA projections as a batched qmv with one 64-thread threadgroup per 8 rows of one head: 4096 tiny threadgroups for embed_q (64 heads x 512 rows, K 256) and 2048 for unembed_out (64 x 256, K 512). In the decode timeline they moved their 5.8-8.9 MB of weights at 0.47-0.65 TB/s, on the dependent chain of every sparse-attention layer. mla_head_qmv keeps each row's arithmetic (qmv_fast's lane mapping, or qmv's tail block when K is exactly one qmv block, as for the 5-bit embed_q) and puts 8 simdgroups x 4 rows in a threadgroup. Shapes MLX routes to qmv_quad (K 64/128 at 4/8 bits) and 2/3-bit weights keep the reference call. Tests: bitwise against QuantizedMultiLinear for 4/5/6/8 bits, K 256/512/1024, group sizes 32/64, and the small-model decode and verify logits with quantized MLA projections. Measured (dependent chain, 64 heads, 8-bit): embed_q 27.6 -> 15.1 us, unembed_out 25.5 -> 12.6 us per call. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): gather one-token sparse-attention rows in one dispatch For one token (no batch mask) the sparse attention clipped the selected indices (maximum, minimum), gathered the latent rows (take_along_axis) and built the selection mask (>= 0): four dependent dispatches on every sparse layer once the context exceeds index_topk. dsa_gather_selected copies the clamped rows (16-byte loads) and writes the mask in one. Pure data movement, so the attention inputs are identical. Tests: the kernel against the reference ops (context 1/300/4099, 1/7/2051 selected rows) and the small model's decode/verify logits with the production kernel family set (fused latent attention off), bitwise. Only one-token decode past index_topk (4k context here) takes this path; see the HC commit for the combined in-model numbers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): run one-token embed_q beside the indexer selection With index selection (context above index_topk), a one-token sparse layer computed embed_q(q) only after the indexer's scores, top-k, index expansion and row gather, although it needs nothing but q. It is now computed right after the query projections and the indexer's query input depends on it without reading it, so it is scheduled first and runs beside the selection chain. Scheduling only: the same kernels on the same inputs (small-model decode/verify logits stay bitwise). Raw pipelined decode (stock wheel, forward and reverse runs): 16.72 -> 16.68 ms/step at 1k context, 17.62 -> 17.58 ms/step at 4k. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): one-token router sigmoid matches the eager kernel on release wheels group_expert_select takes the sigmoid of the fp32 router logits with MLX's eager Sigmoid kernel (it is not fused into a compiled kernel there). On the release wheels that kernel is precompiled with -fno-fast-math and uses the precise exp, while the fused router kernels used the runtime compiled exp: a last-bit difference for a few percent of logits that, after the top-k normalization, changed ~5% of the routing weights by one ulp. The decode then left the reference path at the first decode step (stock wheel, 4096-token context: max logit difference 0.44, tokens diverging at step 11; source builds were exact). This predates perf/glm-decode-2 (perf/glm-decode's one-token and verify routers). The router logits kernels now take the precise sigmoid where eager_sigmoid_precise(float32) says the eager kernel does (probed once, eagerly, from the decoder layer before its compiled FFN block is traced; undecided means the reference router). Test: the one-token router against group_expert_select over 96 draws per shape, including the checkpoint's router bias level (the old test tolerated 2 ulp on one shape and passed by chance with 12 draws on the other). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): keep the one-token shared expert in the routed gate/up kernel Running the shared expert's gate/up as its own dispatch beside the router logits (moe_shared_split, 7db891d5) was meant to hide the router behind it. In the model that dispatch streams its 18 MB at ~0.65 TB/s next to the router and the HC post kernel and the routed kernel then starts after it, so the one-launch kernel (routes selected in-kernel, shared expert as the last slot) is faster. Raw pipelined decode on the stock wheel, same tree with only the family switched (ms/step, forward and reverse order): split 16.61 / 17.58-17.66 (1k / 4k context), one launch 16.35-16.44 / 17.25-17.34. moe_shared_split joins the default-off families; its exact kernels stay covered by the tests (which enable every family). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): keep layer weights out of the compiled FFN trace MLX 0.32.2's mx.compile leaks every multi-output primitive that is an intermediate of a traced graph: compile_dfs rewires the trace by assigning over the old outputs' handles, and array assignment skips the sibling-cycle break that ~array runs (fixed upstream by ml-explore/mlx#4453, not in 0.32.2), so the group's descriptors keep each other alive together with every array they reference. The one-token FFN half that Glm5NextDecoderLayer compiles runs through multi-output fused kernels whose inputs are the layer's weights: since 4a9d32e the route-selecting gate/up kernel (act, indices, scores) reads the routed and shared gate/up weights, and the deferred HC chain's pre/post kernels read the connection weights. As trace constants those weights stayed allocated after the model was dropped: each compiled MoE layer could leak its ~2.5 GB of gate/up experts, so unloading GLM-5.3-Flash-oQ4e after a request left 43-77 GB active (dev server: "Settle barrier timed out ... freed=108.53GB", "Emergency reclaim failed ... active_memory=63.85GB"), and the next load was killed for memory. Upstream main's compiled FFN leaks only the small HC mixing weights this way (up to ~1 MB in the test model). compile_ffn_block() compiles the FFN methods with the arrays of the modules they read (ffn_hc, post_attention_layernorm, mlp) passed as inputs: in the trace they are shape-only placeholders, so a leaked group holds no weights. The same arrays are passed on every call, so the graph, the kernels and the values are unchanged (about 0.5 us more per compiled call for the extra inputs). The MTP verify runtime compiles its FFN through it too. GLM-5.3-Flash-oQ4e on the combined dev tree (e2105bf3 + this commit), engine path in one process, engine_pool-style unload, active memory left: admin-bench requests (pp 32/1024/4096/16384/65536, tg 128): 56.19 GB -> 0.00 GB 1024-token prompt + 2x255 decode tokens + 16k prompt: 58.74 GB -> 1352 B stock mlx 0.32.2 wheel, 1024-token prompt + decode: 61.30 GB -> 1336 B main + perf/glm-decode-2 alone (1024 + 32 tokens): 56.19 GB (main: 0.01 GB) Greedy decode at 1024 tokens 61.2/61.3 -> 61.7/61.8 tok/s, 16k 55.7 -> 56.0 tok/s, identical output. Small GLM-5.3 model of the decode-kernel tests (prefill + 3 one-token steps, 3 runs): 73-220 MB left -> 0-36 bytes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: jundot <jundot@users.noreply.github.com>

…ernels (#4026) * perf(glm5_next): fuse decode/verify op chains into exact kernels Single-token decode and short verify forwards (B=1, L<=8) of GLM-5.3 ran ~6600 dispatches per token: MLX encode on the host (~3.8 us/dispatch) and the GPU were both ~22 ms per step. Replace the longest op chains with JIT kernels that reproduce the reference arithmetic bit for bit: - HC: hc_mix (fp32 RMS + mix GEMV, gemv reduction order) + exact_hc_norm replace cast/RMSNorm/GEMV/sinkhorn/RMSNorm; one-token hc_expand issues the same NAX relaxed-precision matmul2d as MLX's fp32 GEMM (only when MLX itself would: NAX available and MLX_ENABLE_TF32 on) plus the epilogue. - MoE: one-token router (gemv + sigmoid/bias, stable-sort top-k), fused routed+shared gate/up SwiGLU and down+weighted-combine kernels for route counts SwitchGLU leaves unsorted (L*topk < 64). - KDA: one kernel per layer between in_proj and o_proj (conv, SiLU, l2norm, qmv_quad gate projections, safe gate, delta rule, RMSNormGated). - DSA indexer: decode score kernel equal to the Steel tile + mask, and the top-k index expansion (fixes an unusable .strides access). Exactness: 42 unit tests compare every path bitwise with the reference op graph, including an end-to-end small model (decode and verify widths 1..8, with and without TF32 NAX). Engine greedy decode, 256 tokens at pp200/1024/4096: token-identical to OMLX_GLM5_DECODE_FUSION=0. Measured (M5 Ultra, engine greedy, same tree, fusion off -> on): 38.9/38.7/32.7 -> 51.7/51.2/49.0 tok/s at pp200/1024/4096. Raw forward at ctx 200/4096: L=1 24.8/29.8 -> 18.9/20.3 ms, L=4 39.2/45.8 -> 32.3/35.7 ms, L=8 50.3/59.3 -> 42.5/48.8 ms. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): exact verify-block router and wider small decode kernels With the op chains fused, one-token decode is GPU-bound (the host waits on in-flight work ~half the time) and the remaining GPU cost is dominated by weight streaming plus a few latency-bound small kernels. Make those use more of the GPU and fuse the verify router, without changing a bit: - hc_mix: one mix row per 256-thread threadgroup (24 threadgroups per token instead of 6), replaying the 1024-thread rms_looped reduction tree. - one-token router logits: one expert row per simdgroup (the gemv's per-row arithmetic does not depend on rows per thread). - router select: rank-based top-k (one thread per expert, stable-sort tie order, NaNs last) instead of eight serial selection rounds. - verify blocks (2..8 tokens): the router logits reproduce MLX's NAX split-K fp32 GEMM (same relaxed matmul2d chain per K partition, same partition sum), followed by the fused sigmoid/bias/top-k/normalize. Used only where MLX itself runs fp32 GEMMs on NAX with TF32 on. 44 bitwise tests (TF32 NAX paths checked in a TF32-enabled subprocess, including an end-to-end small model with the verify router engaged). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): fold the verify-block shared expert into the MoE gate/up kernel For 2..8-token verify blocks the shared expert's gate and up projections ran as two MLX qmv_wide matmuls plus five clamp/SiLU/multiply kernels per MoE layer. The fused gate/up kernel gets one extra slot that replays qmv_wide's multi-row arithmetic (8 lanes per row, per-group dequantize in 8-value sub-chunks, 4/2/1 shuffle ladder) for all tokens at once and applies the clamped SwiGLU, so the shared weights are read once per layer and seven dispatches become none. The shared down projection keeps MLX's own qmv_wide call. Bit-identical (MoE tests at L=2/4/7 with 8- and 4-bit shared experts assert the new slot is used). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): exact fused latent attention for DSA decode and dense verify GLM-5.3's NoPE MLA attends in the 512-wide latent space, which mx.fast.scaled_dot_product_attention serves with its op fallback: a bf16 scale multiply, a NAX GEMM for the scores with the 64 heads folded into M, a where() mask, a precise softmax, a second NAX GEMM (split along K past ~1k keys) and, for one-token sparse decode, a materialized gather of the 2051 selected latent rows. The small-M GEMMs run on very few threadgroups. Three kernels replay that exactly: the same bf16 16x32x16 relaxed matmul2d chains per K partition (steel_matmul's split-K routing and partition sums, two 16-wide steps per 32-wide k step, zero-padded tails), the softmax_single_row reduction tree, and the same rounding points. Sparse decode reads the selected (clamped) rows in place. Used for one-token decode (dense and sparse) and dense 2..8-token verify blocks with a bool or causal mask when NAX is available; sparse verify keeps _gathered_attention. Tests: 46-case bitwise sweep against SDPA (1..3000 keys, both split-K regimes, sparse gathers with negative indices, L=2..8 masks), plus the end-to-end small models (dense and sparse prompts, TF32 on and off). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): exact sparse latent attention for DSA verify blocks Past 2048 keys a 2..8-token verify block runs _gathered_attention: each token gathers its 2051 selected latent rows ([L, 2051, 512] copies) and the SDPA fallback solves L * 64 one-row problems with MLX's batched gemv (scores) and gemv_t (values), re-reading every token's rows once per head. Replay those reductions exactly (gemv: 4 keys per simdgroup, 16 products per lane, 16/8/4/2/1 ladder; gemv_t: 8 key groups x 4 lanes, 4 simdgroups per 64 columns, leftover keys, 16/8/4 ladder; the shared precise softmax) while reading each selected row in place, once per chunk of 8 heads. Tests: 12 bitwise cases against the gathered SDPA (L = 2..8, 64..2051 selected keys, negative/clamped indices) and the sparse end-to-end small model now exercising the path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): one dispatch per input for the DSA attention projections Decode/verify blocks (B == 1, L <= 8) of the DSA layers now compute the projections that share an input in one multi_qmv dispatch per quantization: q_a, kv_a, the indexer's wk (and weights_proj when the indexer selects) from x, and q_b (with the indexer's wq_b) from qr. multi_qmv replicates qmv_fast for one token and qmv_wide for 2..8 rows per output row, so every projection is bitwise the separate matmul; the indexer takes the precomputed projections. Layers whose projections do not share (bits, group size) keep the reference calls. The one-token router selection goes back to the simdgroup top-k rounds (11.5 us vs 16.2 us for the per-expert rank kernel on GLM-5.3) and now orders NaN scores after every number, lowest index first, like argpartition's sort. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * chore(glm5_next): per-family switch for the fused decode kernels OMLX_GLM5_DECODE_DISABLE (comma-separated STATS keys) or decode_kernels.DISABLED at runtime sends the named families to the reference path, so each kernel's contribution can be measured in one process on the same weights. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): prefetch the latent attention NAX operands The latent scores/values kernels waited on every operand load before each 16x32x16 NAX op (and, for sparse decode, on an index load before the value row load), so a 2048-key partition took ~130 dependent memory round trips. The operands of the next 8 chunks are now loaded ahead of the accumulator chain (the chain and every op are unchanged), selected value rows are staged in threadgroup memory, and the key rows are addressed by offsets instead of an array of device pointers. The Metal compiler produced wrong code for a 4-deep ring in some split-K configurations, so each compile-time configuration is checked bitwise against the SDPA fallback on its first call and falls back to the reference if it differs; tests cover every configuration GLM-5.3 decode and verify use, and several ring depths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): prefetch the verify router's NAX operands Same treatment as the latent kernels: the split-K router logits chain (128 dependent 16x32x16 TF32 ops per partition) now has the operands of the next 8 chunks in flight, and rows are addressed by offsets. Each configuration is checked against group_expert_select on its first call and falls back to the reference if it differs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): route-slot-major MoE grids for verify blocks Verify blocks route 2..8 tokens to overlapping experts (about 24 distinct of 32 routes at 4 tokens, 40 of 64 at 8), but the gate/up grid ran one route slot after another (and the down grid one token after another), so an expert shared by two tokens was streamed from DRAM twice, hundreds of microseconds apart. For blocks of 2+ tokens the slots (tokens) now vary fastest in the grid, so shared experts' row blocks are read back to back. Every threadgroup computes exactly what it did before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): fold the verify-block shared expert down into the combine For 2..8-token blocks the shared expert's down projection ran as its own [T, 2048] x [4096, 2048] quantized matmul before the routed combine kernel added it. The combine kernel now computes it in place for each (token, row) with the multi-row qmv_wide arithmetic that matmul uses (8 lanes per row, per-token accumulators independent of the other tokens), then adds it exactly as before: one dispatch and one [T, 4096] round trip less per MoE layer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): keep latent attention off MLX's gemv-routed score shapes The fused latent attention replicates the SDPA fallback's NAX GEMMs, but MLX computes the scores q @ k.T with gemv_wide when the folded query has 2..15 rows (heads x tokens) and with gemv for a single key, so with a few-head model (the 8-head test model at one token) a handful of outputs differed by one bf16 ULP on some inputs (seed 48 / 300-token prompt, and the combined build's KDA-prefill activations). GLM-5.3 (64 heads, >= 64 rows) never takes those routes; the kernel now declines them. Tests: the small model uses 16 heads (so the one-token sparse path stays covered), a seed x prompt-length x heads sweep (including the failing seed 48 / 300 / 8 heads case), explicit declines, and a KDA-fast-prefill then fused-decode check that runs when that prefill is in the build. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): fused decode experts also read a fused gate_up_proj With the MoE gate/up fusion (perf/moe-fused-gate-up) GLM-5.3's SwitchGLU holds one gate_up_proj ([E, 2N, K]: gate rows then up rows per expert) instead of gate_proj/up_proj, and _decode_experts failed on the missing attributes. The gate/up kernel now takes the expert row stride and the up-row offset as template constants, so the fused tensor is read in place (passed as both operands, up rows N further on) with the same per-row arithmetic; the sort threshold uses the projections SwitchGLU itself passes. Tests: fused layout vs split layout vs reference, bitwise, and vs the fused SwitchGLU forward where the build has it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): leave the fused latent attention off by default In the model the latent attention replicas (one-token / dense verify NAX kernels and the sparse verify rows kernel) cost more than the reference ops they replace: with one call in flight each output tile is a chain of up to 128 dependent 16x32x16 NAX ops at low occupancy (single-call latency e.g. ~300 us vs ~60 us for the 2051-key one-token reference, ~830 us vs ~260 us for a 4-token sparse verify; disabling them made the in-model one-token step 1.5-2.9 ms faster at 1k-4k context and the 4/8-token verify 3.7/6.7 ms faster at 4k). They stay exact and tested, and OMLX_GLM5_DECODE_DISABLE (which now replaces the default-off set) re-enables them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): skip the first-use kernel checks inside mx.compile The latent attention and verify-router first-use checks evaluate a comparison with .item(), which raises while tracing (mx.compile / vmap). When that happens the call now takes the reference path without caching a verdict, and the check runs on the next eager call. (The verify router only runs eagerly in the model; the compiled FFN covers one-token blocks, whose router has no check.) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(glm5_next): fused decode after upstream's fused KDA prefill The small model's 512-token prompt chunks now run #3944's fused KDA prefill prework and norm-gate (glm53_kda_prework); the fused decode and verify paths stay bitwise equal to the reference decode on the caches it writes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(glm5_next): router tests pass on the stock mlx 0.32.2 wheel Two router tests only passed on a source build of mlx. On the stock 0.32.2 wheel (what oMLX ships): - the one-token router stays bitwise at GLM-5.3's shape (288 experts x 4096), but for the small 64 x 512 test shape the wheel's precompiled gemv uses another reduction variant than the one the kernel replays; that shape now allows the summation-order difference (<= 2 fp32 ulp, identical expert indices); - with TF32 NAX matmuls on, the verify-block router's first-use check declines one of the five row counts (its TF32 GEMM rounds differently there) and that configuration keeps the reference path, as designed. The test now accepts a declined configuration; every call is still checked bitwise against the reference. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(glm5_next): skip model-level fused-decode checks under the MTP runtime glm5_next_vlm_runtime.apply() (MTP checkpoints) is process-wide and replaces the layer __call__ methods, which do not route the fused decode kernels. Tests that ran it earlier in the process (e.g. the oQ round trip with nextn weights) made seven model-level checks here fail when they ran first; those checks now skip in that case. Kernel-level tests are unaffected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): fused decode kernels default to M5 (NAX) GPUs The fused decode/verify kernels replay the reference op graph bit for bit on M5 GPUs, where they were validated. On other GPUs the reference ops run different MLX kernels, so the replay is not bitwise, and several kernels use 1024-thread threadgroups that some devices do not allow (the GitHub macOS runners' virtual GPU caps them at 768 or 640 and failed the CI). OMLX_GLM5_DECODE_FUSION now defaults to on only when NAX is available (=1 forces the kernels on, =0 turns them off everywhere); the bitwise kernel tests skip on non-NAX devices. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): KDA decode sigmoids match the eager kernels on release wheels The fused KDA decode kernel evaluates beta = sigmoid(b) and the RMSNormGated gate sigmoid with MLX's Sigmoid expression, whose metal::exp is compiled with the runtime default math. The reference computes both with eager mx.sigmoid calls, and the release mlx wheels run those from the precompiled metallib, which is built with -fno-fast-math (precise exp). On the stock mlx 0.32.2 wheel the two differ in ~20% of sigmoid values; after bf16 rounding the kernel's outputs differed from the reference in a few elements for 6 of 10 test seeds (then in the recurrent state from the next verify block on). Source builds that JIT every kernel use the default exp for mx.sigmoid too, which is why the existing checks passed there. eager_sigmoid_precise() now decides once per dtype which expression the eager kernel reproduces (sweep of 2^17 inputs; None when neither, and the fused path then declines), and the kernel selects exp for those two sites accordingly. Every other exp in the fused kernels has a compiled-graph or custom-kernel reference (nn.silu, compute_g_safe, group_expert_select, the HC kernels) and keeps the runtime-compiled expression. New tests: the probe reproduces mx.sigmoid, and a 10-seed decode/verify sweep of the KDA layer is bitwise equal to the reference (6 seeds failed on the stock wheel before this change). 168 passed on the stock wheel and on the source build. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 98cc9ba245ec6fe4cae7fceed87e33ec45bff33f) * fix(glm5_next): one-token router sigmoid matches the eager kernel on release wheels group_expert_select takes the sigmoid of the fp32 router logits with MLX's eager Sigmoid kernel (it is not fused into a compiled kernel there). On the release wheels that kernel is precompiled with -fno-fast-math and uses the precise exp, while the fused router kernels used the runtime compiled exp: a last-bit difference for a few percent of logits that, after the top-k normalization, changed ~5% of the routing weights by one ulp. The decode then left the reference path at the first decode step (stock wheel, 4096-token context: max logit difference 0.44, tokens diverging at step 11; source builds were exact). This predates perf/glm-decode-2 (perf/glm-decode's one-token and verify routers). The router logits kernels now take the precise sigmoid where eager_sigmoid_precise(float32) says the eager kernel does (probed once, eagerly, from the decoder layer before its compiled FFN block is traced; undecided means the reference router). Test: the one-token router against group_expert_select over 96 draws per shape, including the checkpoint's router bias level (the old test tolerated 2 ulp on one shape and passed by chance with 12 draws on the other). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 204d8277950adb3d2546693dca6329cc31b5fc34) * perf(glm5_next): start evaluating one-token decode steps early A GLM-5.3 decode step is ~800 dependent dispatches and its Python graph build takes ~2.7 ms. mlx allows ~10 command buffers in flight, and with its default per-buffer size budget (every expert or projection weight input counts in full) one step commits ~170 buffers, so while the next step's graph is being built the GPU drains the few buffers still queued and then idles: ~1.3 ms per token at 1k context, ~1.6 ms at 4k (devbuild timeline: "host-starved gaps"). One-token forwards now call mx.async_eval on the hidden state every 8 layers (OMLX_GLM5_DECODE_EVAL_EVERY, 0 = off), so the first layers are encoded and running while the later ones are still being built. The computed values are unchanged (scheduling only); the MTP runtime's replacement forward follows the same policy. Raw decode steps, GLM-5.3-Flash-oQ4e, M5 Ultra, mlx's default buffer budget: 1k context 17.99 -> 16.83 ms/token (55.6 -> 59.4 tok/s), 4k 19.45 -> 17.83 ms (51.4 -> 56.1). With an op-count-only buffer budget (no host gap to hide) it is neutral: 16.72 -> 16.67, 17.77 -> 17.71. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): fused KDA decode for layers with mixed projection bits The fused KDA decode step needs the q|k|v|f_a|g_a|b projection of the layer as one row block. It was only built when all six projections share one quantization, so a layer whose v_proj is 5-bit among 8-bit q/k/gates (GLM-5.3-Flash-oQ4e layer 40) decoded through the reference body: six projection matmuls, concatenations, the short conv, SiLU, l2norms, gate projections, the gated-delta kernel and RMSNormGated as ~35 dependent dispatches instead of three. Such layers now get one fused weight per quantization (a single-member group uses the layer's own arrays) and the decode projection runs one matmul per group, putting the rows back in q..b order with one concatenation of contiguous runs. Each row is the row of its own layer's matmul (same qmv kernel and per-row arithmetic), so the block equals the reference projections and the fused kernel sees the same input. Test: a KDA layer with a 5-bit or 4-bit v_proj decodes and verifies bitwise like the reference and takes the fused kernel. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): 5-bit KDA gate projections inside the decode kernel The fused KDA decode kernel replays the forget/output gate projections (f_b, g_b: K = 128 rows) for 4- and 8-bit weights, which MLX runs as qmv_quad. GLM-5.3-Flash-oQ4e quantizes them to 5 bits in 8 of its 34 KDA layers; there the two projections ran as separate quantized matmuls before the kernel: MLX's one-row qmv with only 16 of 32 lanes active for K = 128, two latency-bound dependent dispatches (~20 us each in a one-dispatch-per- buffer timeline) per layer and token. For one token the kernel now computes these rows itself with qmv's arithmetic: lanes 0..15 load 8 values each (load_vector_safe/qdot_safe with N = 8, identical to load_vector/qdot), the other lanes contribute zero, one simd_sum per row. Verify blocks (2+ rows go to qmv_wide) keep the precomputed projections. Tests: the 10-seed decode/verify sweep now also runs with 5-bit gates and checks that one-token steps take the in-kernel rows; 322 passed on the stock wheel and on the source build. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): select one-token MoE routes inside the gate/up kernel A one-token MoE block ran the router as two dependent dispatches (logits GEMV, then a single-simdgroup top-k select) before the gate/up kernel. The routed gate/up threadgroups now replay the select kernel's loop on the biased sigmoid scores themselves (the same code, moved into a shared helper: stable descending order, ties to the lower expert index, NaNs last), and slot 0 / tile 0 publishes the routes and normalized routing weights for the down/combine kernel. The select dispatch disappears; the selection is ~1 us of simdgroup work at the start of each routed threadgroup. Verify blocks keep moe_router_rows + select. OMLX_GLM5_DECODE_DISABLE= router_select_fused restores the two-kernel router. Test: routes and outputs bitwise equal to the router select path and to the reference MoE for GLM-5.3's shape (288 experts, K 4096), including exact score ties. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): overlap the one-token router with the shared expert In a one-token MoE block the shared expert's gate/up does not depend on the router, but it ran inside the routed gate/up kernel, after the router logits kernel: a 4.7 MB fp32 GEMV spread over 72 threadgroups that is mostly latency. The shared expert's gate/up now runs as its own dispatch (the shared slot of the fused kernel, same code and arithmetic) ahead of the routed one. Neither it nor the router logits kernel reads the other's output, so mlx encodes no barrier between them and they run concurrently: the router's latency hides under the shared expert's 17.8 MB weight stream. The routed gate/up (with the in-kernel top-k) takes the shared activation as an ordering-only input, and the down/combine kernel reads it from its own buffer (moe_down_combine(shared_act_sep=...), same qmv arithmetic as the in-act slot). OMLX_GLM5_DECODE_DISABLE=moe_shared_split restores the single gate/up launch. Test: bitwise equal to the single-launch path, the router-select path and the reference MoE for GLM-5.3's shape, including exact score ties. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(glm5_next): model-level checks require the one-token router paths The small-model decode/verify logits checks (bitwise against the reference model, dense and sparse DSA attention, seed sweep) now also assert that the in-kernel route selection and the separate shared-expert gate/up launch were taken. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): one-token dense MLP gate/up and SwiGLU in one dispatch GLM-5.3's first three layers use a dense MLP (8-bit, 4096 -> 12288). For one token its gate and up projections ran as two quantized matmuls (concurrent, ~0.89 TB/s together for 100 MB in a one-dispatch-per-buffer timeline) followed by the clamped SwiGLU as a third dispatch. They now run as one dispatch of the shared-expert gate/up kernel (Glm5NextMLP is the shared expert's module; same one-row qmv arithmetic per row and the same clamped-SwiGLU epilogue), then the reference down projection. Test: bitwise equal to the eager and the compiled reference MLP (8- and 4-bit). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): bitonic top-k for decode/verify indexer rows Past index_topk tokens of context the DSA indexer picks 512 of up to 2048 pooled blocks per decode/verify row with the native radix-select kernel: two histogram passes whose threadgroup atomics contend on the few bins that hold similar scores, and two 256-bin scans by a single thread (~15-20 us per sparse layer in a one-dispatch-per-buffer timeline). Its output is a function of the scores alone: keys are the 16-bit order-preserving score bits, strictly greater keys fill the output in index order, then threshold ties in index order. dsa_topk_rows finds the threshold with a threadgroup bitonic sort of (key, index) pairs (one 1024-thread threadgroup per row, P <= 2048) and emits the same output with the native kernel's raking index-order pass. Rows beyond 2048 blocks and other select widths keep the native kernel. Tests: identical indices and order to the native kernel for 512..2048 blocks, 1/4/8 rows, exact ties, -1e30 masked blocks and a NaN; the small-model sparse decode/verify logits check requires this path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): load indexer query fragments ahead of the score MMAs dsa_decode_scores loaded each 8-wide query fragment inside the per-head chain of 16 accumulator-dependent simdgroup MMAs, so every step waited on its own load. Each head's fragments are now loaded before its chain; the MMA sequence (and so every score) is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): fold each one-token HC expand into the next HC pre A one-token half-layer ran three dependent, latency-bound HC dispatches around its branch: hc_expand of the previous branch output, hc_mix (RMS + 24-row mix GEMV) and exact_hc_norm (sinkhorn, collapse, RMSNorm). The branch input only needs the four pre rows of the mix, so: - hc_pre_fused computes the mix rows four per 1024-thread threadgroup (hc_mix's multi-row layout) and threadgroup 0, which owns the pre rows, finishes exact_hc_norm's collapse and RMSNorm itself: one dispatch on the chain instead of two; - hc_post_mm (post weights, sinkhorn comb and hc_expand_one's NAX comb product of the residual) only feeds the next expand; the branch input depends on it without binding it, so it is scheduled first and runs beside the branch's first kernel; - the next hc_pre_fused applies hc_expand_one's epilogue to the branch output while loading h and stores h (the next residual), so the expand dispatch disappears. The model loops (vendor and MTP runtime) pass the pending expand between layers and materialize the last one. Every value keeps the reference kernels' arithmetic and order (the sinkhorn code is compiled ahead of the MPP include, as in exact_hc_norm). Tests: a chained kernel check against HyperConnection + RMSNorm + hc_expand (hidden 4096/2048/1024, TF32 NAX on) and the small-model decode/verify logits, also under the MTP runtime loop, bitwise. Measured (stock wheel): in a dependent chain of 90 HC pres alone, 26.0 -> 20.6 us per half-layer; with a two-qmv branch per half-layer (71 MB of 8-bit weights) 103.6 -> 100.0 us. Raw pipelined decode of the model with this commit and the next two (MLA head qmv, sparse row gather): 16.81 -> 16.72 ms/step at 1k context, 17.94 -> 17.62 ms/step at 4k. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): one-token MLA embed_q/unembed_out in wider threadgroups QuantizedMultiLinear runs the per-head MLA projections as a batched qmv with one 64-thread threadgroup per 8 rows of one head: 4096 tiny threadgroups for embed_q (64 heads x 512 rows, K 256) and 2048 for unembed_out (64 x 256, K 512). In the decode timeline they moved their 5.8-8.9 MB of weights at 0.47-0.65 TB/s, on the dependent chain of every sparse-attention layer. mla_head_qmv keeps each row's arithmetic (qmv_fast's lane mapping, or qmv's tail block when K is exactly one qmv block, as for the 5-bit embed_q) and puts 8 simdgroups x 4 rows in a threadgroup. Shapes MLX routes to qmv_quad (K 64/128 at 4/8 bits) and 2/3-bit weights keep the reference call. Tests: bitwise against QuantizedMultiLinear for 4/5/6/8 bits, K 256/512/1024, group sizes 32/64, and the small-model decode and verify logits with quantized MLA projections. Measured (dependent chain, 64 heads, 8-bit): embed_q 27.6 -> 15.1 us, unembed_out 25.5 -> 12.6 us per call. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): gather one-token sparse-attention rows in one dispatch For one token (no batch mask) the sparse attention clipped the selected indices (maximum, minimum), gathered the latent rows (take_along_axis) and built the selection mask (>= 0): four dependent dispatches on every sparse layer once the context exceeds index_topk. dsa_gather_selected copies the clamped rows (16-byte loads) and writes the mask in one. Pure data movement, so the attention inputs are identical. Tests: the kernel against the reference ops (context 1/300/4099, 1/7/2051 selected rows) and the small model's decode/verify logits with the production kernel family set (fused latent attention off), bitwise. Only one-token decode past index_topk (4k context here) takes this path; see the HC commit for the combined in-model numbers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): run one-token embed_q beside the indexer selection With index selection (context above index_topk), a one-token sparse layer computed embed_q(q) only after the indexer's scores, top-k, index expansion and row gather, although it needs nothing but q. It is now computed right after the query projections and the indexer's query input depends on it without reading it, so it is scheduled first and runs beside the selection chain. Scheduling only: the same kernels on the same inputs (small-model decode/verify logits stay bitwise). Raw pipelined decode (stock wheel, forward and reverse runs): 16.72 -> 16.68 ms/step at 1k context, 17.62 -> 17.58 ms/step at 4k. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): one-token router sigmoid matches the eager kernel on release wheels group_expert_select takes the sigmoid of the fp32 router logits with MLX's eager Sigmoid kernel (it is not fused into a compiled kernel there). On the release wheels that kernel is precompiled with -fno-fast-math and uses the precise exp, while the fused router kernels used the runtime compiled exp: a last-bit difference for a few percent of logits that, after the top-k normalization, changed ~5% of the routing weights by one ulp. The decode then left the reference path at the first decode step (stock wheel, 4096-token context: max logit difference 0.44, tokens diverging at step 11; source builds were exact). This predates perf/glm-decode-2 (perf/glm-decode's one-token and verify routers). The router logits kernels now take the precise sigmoid where eager_sigmoid_precise(float32) says the eager kernel does (probed once, eagerly, from the decoder layer before its compiled FFN block is traced; undecided means the reference router). Test: the one-token router against group_expert_select over 96 draws per shape, including the checkpoint's router bias level (the old test tolerated 2 ulp on one shape and passed by chance with 12 draws on the other). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): keep the one-token shared expert in the routed gate/up kernel Running the shared expert's gate/up as its own dispatch beside the router logits (moe_shared_split, 7db891d5) was meant to hide the router behind it. In the model that dispatch streams its 18 MB at ~0.65 TB/s next to the router and the HC post kernel and the routed kernel then starts after it, so the one-launch kernel (routes selected in-kernel, shared expert as the last slot) is faster. Raw pipelined decode on the stock wheel, same tree with only the family switched (ms/step, forward and reverse order): split 16.61 / 17.58-17.66 (1k / 4k context), one launch 16.35-16.44 / 17.25-17.34. moe_shared_split joins the default-off families; its exact kernels stay covered by the tests (which enable every family). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(glm5_next): keep layer weights out of the compiled FFN trace MLX 0.32.2's mx.compile leaks every multi-output primitive that is an intermediate of a traced graph: compile_dfs rewires the trace by assigning over the old outputs' handles, and array assignment skips the sibling-cycle break that ~array runs (fixed upstream by ml-explore/mlx#4453, not in 0.32.2), so the group's descriptors keep each other alive together with every array they reference. The one-token FFN half that Glm5NextDecoderLayer compiles runs through multi-output fused kernels whose inputs are the layer's weights: since 4a9d32e the route-selecting gate/up kernel (act, indices, scores) reads the routed and shared gate/up weights, and the deferred HC chain's pre/post kernels read the connection weights. As trace constants those weights stayed allocated after the model was dropped: each compiled MoE layer could leak its ~2.5 GB of gate/up experts, so unloading GLM-5.3-Flash-oQ4e after a request left 43-77 GB active (dev server: "Settle barrier timed out ... freed=108.53GB", "Emergency reclaim failed ... active_memory=63.85GB"), and the next load was killed for memory. Upstream main's compiled FFN leaks only the small HC mixing weights this way (up to ~1 MB in the test model). compile_ffn_block() compiles the FFN methods with the arrays of the modules they read (ffn_hc, post_attention_layernorm, mlp) passed as inputs: in the trace they are shape-only placeholders, so a leaked group holds no weights. The same arrays are passed on every call, so the graph, the kernels and the values are unchanged (about 0.5 us more per compiled call for the extra inputs). The MTP verify runtime compiles its FFN through it too. GLM-5.3-Flash-oQ4e on the combined dev tree (e2105bf3 + this commit), engine path in one process, engine_pool-style unload, active memory left: admin-bench requests (pp 32/1024/4096/16384/65536, tg 128): 56.19 GB -> 0.00 GB 1024-token prompt + 2x255 decode tokens + 16k prompt: 58.74 GB -> 1352 B stock mlx 0.32.2 wheel, 1024-token prompt + decode: 61.30 GB -> 1336 B main + perf/glm-decode-2 alone (1024 + 32 tokens): 56.19 GB (main: 0.01 GB) Greedy decode at 1024 tokens 61.2/61.3 -> 61.7/61.8 tok/s, 16k 55.7 -> 56.0 tok/s, identical output. Small GLM-5.3 model of the decode-kernel tests (prefill + 3 one-token steps, 3 runs): 73-220 MB left -> 0-36 bytes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): run Lightning MTP verify forwards on the fused decode kernels With the glm5_next MTP runtime applied, a verify forward (gdn_sink set) ran the runtime's own layer bodies: the attention half-layer used the stock HC ops and the reference KDA body (short conv, SiLU, l2norm, gate projections, delta rule and RMSNormGated as separate ops), so MTP verify missed the exact fused kernels plain decode and DFlash verify use. The runtime also replaces Glm5NextLinearAttention.__call__ process-wide with a body that never tried the fused decode step, so once it was applied (any MTP-capable GLM-5.3 checkpoint, MTP on or off) plain decode ran the reference KDA body too. - The runtime's KDA call tries the vendor's fused decode/verify step first, for plain decode and verify blocks alike. A verify block records a KdaStepCapture (the kernel's inputs: fused input projection, the a_pre/gate_pre rows when the gate projections run outside the kernel, and the entry conv/recurrent states) instead of the reference tuple. - rollback_speculative_cache replays a partial accept through the same kernel over the kept rows. Each kernel row depends only on earlier rows (causal conv, per-row l2norm and gate projections, sequential delta rule), so the rebuilt states are bit for bit those the block held after that row. A block kept whole needs no replay. Every recurrent state is rebuilt before any cache moves, and a capture of the wrong width is refused. - The verify decoder layer uses _decode_hc_pre / _decode_hc_expand like the vendor call (same values as attn_hc + input_layernorm / hc_expand). - The compiled verify FFN block traced the reference router at every verify width: moe_router_rows' first-use exactness check needs an eager evaluation and declines inside mx.compile, and the trace keeps what it saw. One eager router call per width on a fixed synthetic block settles the check before the first trace (the routes are the same either way). - The batch generator arms the verify-shape qmm routes (qwen35_verify_qmm) around each verify forward; they take QuantizedLinear calls of 3+ rows. A KDA layer whose projections do not share one quantization (GLM-5.3 layer 40, 5-bit v_proj) makes such calls only in the reference body, so while the routes are installed and armed its 3+ row blocks keep that body and verify values never depend on the path. - OMLX_GLM5_MTP_FUSED_VERIFY=0 restores the reference ops in verify forwards. Exactness (tests/test_glm5_next_mtp_fused_verify.py): two identical small quantized GLM-5.3 models with an MTP layer (8-bit, 5-bit and mixed-width KDA projections; dense and index-selecting sparse attention) run MTP verify cycles of widths 1-8 with every accept count 0..depth, fused path vs the previous verify path: logits and hidden states, every cache array and position after the rollback, the MTP head forwards (fused kernels vs reference ops) and the next decode step are bitwise equal, with TF32 NAX on (production default) and off and with the verify-qmm routes armed. A seed sweep ran 1800 such cycles (40 seeds), all bitwise equal. End to end through mlx-lm's BatchGenerator and the MTP loop the fused path emits the reference path's greedy tokens. The decode-kernel model checks no longer skip once the MTP runtime is applied; they pass under it. Measured on GLM-5.3-Flash-oQ4e, M5 Ultra, 1k context, dev stack (m5u/integration), stock mlx 0.32.2, raw forwards, one process per mode, median of 20, build + eval per forward, old verify path -> this commit: verify 1 row 20.37 -> 17.31 ms verify 2 rows 31.98 -> 28.81 ms verify 3 rows 36.54 -> 32.54 ms verify 4 rows 40.62 -> 36.10 ms rollback, block kept whole 1.49 -> 0.64 ms (no replay); partial accept 1.67 -> 1.95-2.2 ms (fused replay) plain decode with the runtime applied 18.57 -> 16.49 ms/token pipelined (53.8 -> 60.6 tok/s) The MTP head forward was already on the fused kernels (fold of 1/2/3 rows 1.42/1.73/1.88 ms). No end-to-end MTP A/B: the local checkpoint has no nextn layer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(glm5_next): start evaluating MTP verify blocks early One-token decode steps start evaluating every _DECODE_EVAL_EVERY (8) layers while the rest of the graph is still being built; verify blocks (2..8 rows) did not, so with mlx's default command-buffer budget the GPU idled through most of each verify forward's ~2.7 ms graph build. The MTP runtime's model loop now applies the same policy to blocks of up to 8 rows. Scheduling only: a test compares logits, hidden states and rolled-back caches bitwise against the lazy forward. OMLX_GLM5_MTP_VERIFY_EARLY_EVAL=0 turns it off. Measured on GLM-5.3-Flash-oQ4e, M5 Ultra, 1k context, dev stack, stock mlx 0.32.2, raw forwards, one process per mode, median of 20, build + eval per verify forward (fused verify without -> with early evaluation): 1 row 17.31 -> 17.27 ms, 2 rows 28.81 -> 26.08, 3 rows 32.54 -> 30.06, 4 rows 36.10 -> 33.39 Against the old verify path (both commits): 20.37 -> 17.27 (-15%), 31.98 -> 26.08 (-18%), 36.54 -> 30.06 (-18%), 40.62 -> 33.39 (-18%) ms. Devbuild GPU busy of a 2-row verify forward + rollback: 25.5 ms, so the fused 2-row verify is now GPU-bound (1.5x a 16.3 ms one-token step). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: jundot <jundot@users.noreply.github.com>

…rify traces Upstream ml-explore#631 declares the fused four-way projection as compile state in the decode and segment traces. The DFlash2 verify traces read it the same way, so they take the same state (traceState(forLayers:)). SiblingCycleTests probes the two mlx causes: an erased compiled function that splits a captured constant keeps it alive (ml-explore/mlx#3932), and assigning over a multi-output sibling skips the cycle break in ~array (ml-explore/mlx#4453). (rebased from 9ad90d7; the decode and segment halves merged upstream as ml-explore#631)

The assignment case is ml-explore/mlx#4453, merged 2026-09-11 and not in the pinned mlx; the compiled split case is ml-explore/mlx#3932, open.

Every model build leaked for good: ~5 MB for Gemma E4B, ~20 MB for Qwen3.5 2B. Idle unloads made a rebuild out of every return to the keyboard, so the leak grew with uptime (~359k MLX array nodes after 38 h). The MLX that mlx-swift 0.32.3 bundles predates ml-explore/mlx#4453: `loadWeights` quantizes each layer's placeholder weights into sibling triples, and the loaded weights then replace them by assignment, which never frees them. - On unload, keepBlueprint records the configuration, the tokenizer directory and the shape of each quantized layer. It captures no arrays, so no weights stay alive. - rebuild makes the model from the registry, puts lazy-zero QuantizedLinear/QuantizedEmbedding placeholders in place, then calls loadWeights: nothing is left to quantize, so nothing leaks. - A blueprint applies only to the same snapshot directory and the same VLM/LLM build; a failed rebuild falls back to a fresh load. Memory pressure drops all blueprints. - The task that reads the blueprint releases the old model before Memory.clearCache, so the freed weights do not refill the pool. - Model and drafter builds run under a scoped random state: from the global state, every placeholder draw grew a lazy key chain that the state held for the app's life. - dropUnusedTowers keeps the real shapes as lazy zeros (never allocated), so the blueprint records them. Rebuild takes 0.44-0.66 s vs 0.75-1.6 s for a fresh load. GPU memory is 0 MB after unload, and outputs are identical. Verified on Gemma E4B (with the VLM fix model and its drafter), Qwen3.5 and MiniCPM5. Qwen3.5 still leaks ~13 MB per reload from a different MLX bug (compile-tape erase); that one needs an MLX update.

* Fix Log and Equal is_equivalent ignoring primitive state (ml-explore#4266) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * Round mxfp8 block scales up to avoid saturation (ml-explore#4353) Co-authored-by: Daniel Hiltgen <daniel.hiltgen@ollama.com> Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * Fix quantized kernels in JIT build (ml-explore#4372) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <katlun@windyviews.com> Co-authored-by: Cheng <zcbenz@gmail.com> * 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 <git@zcbenz.com> * 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 <99216956+pradyotRanjan@users.noreply.github.com> Co-authored-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * Fix mx.from_fp8 E4M3FN NaN decode with branchless carry (ml-explore#4376) Co-authored-by: Saud <saudkamran@Sauds-MacBook-Pro.local> * 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 <git@zcbenz.com> * Validate GGUF tensor dimensions (ml-explore#4378) Co-authored-by: Cheng <git@zcbenz.com> * Fix crash on ellipsis indexing with too many trailing indices (ml-explore#4396) Co-authored-by: Anastasiia Filippova <a_filippova@apple.com> * [CUDA] Cholesky via cuSOLVER (ml-explore#4208) Co-authored-by: Cheng <git@zcbenz.com> * 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 <vaibhavsri1712@gmail.com> * Fix Metal grid sizing for strided scans (ml-explore#4420) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * Fix integer div-by-zero crash in CPU backend (sort, scan, binary ops) (ml-explore#4442) Co-authored-by: Cheng <git@zcbenz.com> * Use 64-bit file seeks on Windows (ml-explore#4456) Co-authored-by: Cheng <git@zcbenz.com> * Use precise::exp in Sigmoid so compiled and eager sigmoid agree (ml-explore#4461) Co-authored-by: Pierre Lamy <pierrelamy427@gmail.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * win: Use DXGI for accurate WDDM VRAM memory budget (ml-explore#4457) Co-authored-by: Cheng <git@zcbenz.com> * 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 <rohan1gautam@gmail.com> * 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 <ranranhaoranzhang@gmail.com> * 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 <19340816+wyanzhao@users.noreply.github.com> Co-authored-by: Cheng <git@zcbenz.com> * Fix jaccl ring all_gather: direction 1 slice was not mirrored (ml-explore#4443) Co-authored-by: Dave Hind <davehind42@gmail.com> Co-authored-by: Anastasiia Filippova <a_filippova@apple.com> * Skip Metal-only gated delta kernel tests on other backends(fix CI) (ml-explore#4522) Co-authored-by: Anastasiia Filippova <a_filippova@apple.com> * [CUDA] Add global scale support to gather_qmm (ml-explore#4507) Co-authored-by: Anastasiia Filippova <a_filippova@apple.com> * [BUG][Metal] Deadlock in fence (ml-explore#4552) Co-authored-by: Cheng <zcbenz@gmail.com> * Fix ordering of complex64 with NaNs (ml-explore#4519) * Improve metal memory usage for SDPA D256 (ml-explore#4505) Co-authored-by: Cheng <git@zcbenz.com> * Add support for matrix transpose (ml-explore#4402) Co-authored-by: Cheng <git@zcbenz.com> * 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 <a_filippova@apple.com> * 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 <git@zcbenz.com> * Raise IndexError for out of bounds axes (ml-explore#4484) * Fix floor_divide for integers (ml-explore#4515) Co-authored-by: Cheng <git@zcbenz.com> * Fix cpu binary ops on large data (ml-explore#4517) Co-authored-by: Cheng <git@zcbenz.com> * Do not throw from CUDA destructors and avoid implicit default streams in eval/compile (ml-explore#4514) Co-authored-by: Cheng <git@zcbenz.com> * [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 <zcbenz@gmail.com> * Add Python 3.15 and 3.15t support (ml-explore#4555) Co-authored-by: Cheng <git@zcbenz.com> * Add fixed size `unique` operation (ml-explore#4501) Co-authored-by: Cheng <zcbenz@gmail.com> * Fix normalizing single-element slice with negative stride (ml-explore#4446) Co-authored-by: Cheng <git@zcbenz.com> * 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 <zcbenz@gmail.com> * Clear CUDA errors once reported and raise eval errors in buffer protocol (ml-explore#4523) Co-authored-by: Cheng <zcbenz@gmail.com> * 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 <99216956+pradyotRanjan@users.noreply.github.com> Co-authored-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Co-authored-by: Cheng <zcbenz@gmail.com> * 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 <zcbenz@gmail.com> * 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 <git@zcbenz.com> * Clamp slice bounds when a slice is combined with an array index (ml-explore#4513) Co-authored-by: Cheng <zcbenz@gmail.com> * Fix grouped Conv1d/Conv2d init to use in_channels // groups (ml-explore#4544) Co-authored-by: Cheng <zcbenz@gmail.com> * Fix mx.view misaligned offset (ml-explore#4585) Co-authored-by: Cheng <git@zcbenz.com> * [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 <zcbenz@gmail.com> * Fix empty unsigned sum/prod failing on Metal (ml-explore#4546) Co-authored-by: Cheng <git@zcbenz.com> * [CUDA] Add gather_qmm_rhs_sm80 (ml-explore#4554) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <60583610+gokay-ai@users.noreply.github.com> Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * Unroll TK=4 keys per iter in sdpa_vector_2pass_1 (ml-explore#4596) Co-authored-by: Cheng <git@zcbenz.com> * Fix CPU float16 SIMD comparison masks (ml-explore#4602) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <gothravenllm@gmail.com> Co-authored-by: Cheng <git@zcbenz.com> * [metal] Add env var based residency refresh interval (ml-explore#4633) * Promote integer inputs in logcumsumexp (ml-explore#4625) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * Fix 32 bit dispatch grid overflow in Metal reductions (ml-explore#4437) Co-authored-by: Cheng <git@zcbenz.com> * Fix gpu boolean masking in 0 dim array (ml-explore#4635) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * Add 1-bit affine quantization support (Metal) (ml-explore#3161) Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> Co-authored-by: bri-prism <288398250+bri-prism@users.noreply.github.com> * 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 <git@zcbenz.com> * [Metal] Fix SDPA max threadgroup for M1/M2 generation (ml-explore#4643) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * Validate .npy shape in mlx/io/load.cpp (ml-explore#4657) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * Fix direct CPU conv writing past output (ml-explore#4668) Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * 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 <git@zcbenz.com> * Fix metal conv bounds checks and integer overflows (ml-explore#4672) Co-authored-by: Cheng <git@zcbenz.com> * python: Make arctan2/radians/degrees support scalar inputs (ml-explore#4679) Co-authored-by: Cheng <git@zcbenz.com> * Fix missing add_temporary in pad_gpu/fill_gpu (ml-explore#4660) Co-authored-by: “Ronald <“Ronaldmannak@me.com”> * 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 <git@zcbenz.com> * 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 <haojin.yang@greenbit.ai> * [CUDA] Fix grouped gemm with large data (ml-explore#4696) Co-authored-by: yanghaojin <haojin.yang@greenbit.ai> Co-authored-by: Cheng <git@zcbenz.com> * 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 <git@zcbenz.com> * Allow configuring the CPU JIT cache directory (ml-explore#4714) Co-authored-by: Cheng <git@zcbenz.com> * Make `broadcast_arrays` return a tuple (ml-explore#4713) Co-authored-by: Cheng <git@zcbenz.com> * Fix unaligned scalar reads and oversized contiguous buffer views (ml-explore#4721) Co-authored-by: Cheng <git@zcbenz.com> * [Metal] Add staged affine_qmm_t_nax kernels with split-k (ml-explore#4682) Signed-off-by: Yanzhao Wang <yzwang1996@gmail.com> Co-authored-by: Cheng <git@zcbenz.com> * [Metal] Add fast route for dequantizing fp4 to half in fp_qmm_t (ml-explore#4709) Co-authored-by: Cheng <git@zcbenz.com> * [Metal] Add qmv_mma/qmv_nax kernels for 6-16 rows with large N/K (ml-explore#4712) Co-authored-by: Cheng <git@zcbenz.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * Fix data race in distributed ring on pending send/recv lists (ml-explore#4737) Co-authored-by: Cheng <git@zcbenz.com> * Use data_size for gather/scatter large-offset checks (ml-explore#4724) Co-authored-by: Cheng <git@zcbenz.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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) <noreply@anthropic.com> * [CUDA] Fix gather_qmm for experts past 2^32 weight elements (ml-explore#4738) Co-authored-by: Cheng <git@zcbenz.com> * 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * 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 <noreply@anthropic.com> --------- Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Signed-off-by: Vaibhav Srivastava <vaibhavsri1712@gmail.com> Signed-off-by: Ranran Haoran Zhang <ranranhaoranzhang@gmail.com> Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com> Signed-off-by: Yanzhao Wang <yzwang1996@gmail.com> Co-authored-by: rohith <kapellirohith@gmail.com> Co-authored-by: Cheng <git@zcbenz.com> Co-authored-by: Vladimir Iglovikov <ternaus@users.noreply.github.com> Co-authored-by: Brian C. <94733710+deBrian07@users.noreply.github.com> Co-authored-by: Cheng <zcbenz@gmail.com> Co-authored-by: Erwin Zhang <59893706+erwinzhang7@users.noreply.github.com> Co-authored-by: Daniel Hiltgen <dhiltgen@users.noreply.github.com> Co-authored-by: Daniel Hiltgen <daniel.hiltgen@ollama.com> Co-authored-by: Aaishwarya Mishra <aaishwarymishra@gmail.com> Co-authored-by: YH Yan <strayberry0w0@gmail.com> Co-authored-by: Dwijen Patel <dwijen@gmail.com> Co-authored-by: Anastasiia Filippova <a_filippova@apple.com> Co-authored-by: Ayaan Gazali <ayaangazali.work@gmail.com> Co-authored-by: katlun-lgtm <katlun@gmail.com> Co-authored-by: katlun-lgtm <katlun@windyviews.com> Co-authored-by: Adityaj0 <93090622+Adityaj0@users.noreply.github.com> Co-authored-by: Eyüp Can Akman <eyupcanakman@gmail.com> Co-authored-by: Pradyot Ranjan <99216956+prady0t@users.noreply.github.com> Co-authored-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com> Co-authored-by: Jagrit Digani <digani@apple.com> Co-authored-by: Philip John Basile <PBasile@Basilecom.com> Co-authored-by: Alessio Pollero <alessio.pollero@ventureone.ae> Co-authored-by: Duhyeon, Kim <49020301+dudududukim@users.noreply.github.com> Co-authored-by: Atirna <noctiselara0@gmail.com> Co-authored-by: Saud Kamran <32072820+saud5150@users.noreply.github.com> Co-authored-by: Saud <saudkamran@Sauds-MacBook-Pro.local> Co-authored-by: YH Yan <yh.yukiyan@gmail.com> Co-authored-by: Roshan Sharma <rs127@fordham.edu> Co-authored-by: sashko-zakharchuk <oleksandr.zakharchuk@gmail.com> Co-authored-by: XXXXRT666 <157766680+XXXXRT666@users.noreply.github.com> Co-authored-by: Vaibhav Srivastava <vaibhavsri1712@gmail.com> Co-authored-by: Robin Gieseke <38696388+TheDarkchip@users.noreply.github.com> Co-authored-by: 2sumtech <2sumtech@gmail.com> Co-authored-by: Vedant Singhal <88609765+Ved235@users.noreply.github.com> Co-authored-by: JasonHonKL <148705846+JasonHonKL@users.noreply.github.com> Co-authored-by: Simeet Nayan <78461155+simeetnayan81@users.noreply.github.com> Co-authored-by: Cdimoy <50119145+Cdimoy@users.noreply.github.com> Co-authored-by: Alessio Pollero <alessio.pollero@gmail.com> Co-authored-by: Rohan Gautam <rohan1gautam@gmail.com> Co-authored-by: michalk8 <46717574+michalk8@users.noreply.github.com> Co-authored-by: Nilesh Patil <128893479+nileshpatil6@users.noreply.github.com> Co-authored-by: Marcos Ashton <59975445+MarcosAsh@users.noreply.github.com> Co-authored-by: pierre427 <pierre@userid.org> Co-authored-by: Pierre Lamy <pierrelamy427@gmail.com> Co-authored-by: Yanzhao Wang <19340816+wyanzhao@users.noreply.github.com> Co-authored-by: David Koski <46639364+davidkoski@users.noreply.github.com> Co-authored-by: Isaac Hernández <isadohergar@gmail.com> Co-authored-by: Tudorica Constantin-Alexandru <tudalex@gmail.com> Co-authored-by: Ritaban Mitra <39273308+Ritabanm@users.noreply.github.com> Co-authored-by: Ranran <ranranhaoranzhang@gmail.com> Co-authored-by: Tommaso Pegolotti <31819389+tpegolotti@users.noreply.github.com> Co-authored-by: Drifter4242 <davehind@yahoo.co.uk> Co-authored-by: Dave Hind <davehind42@gmail.com> Co-authored-by: Valentin Roussellet <vroussellet@apple.com> Co-authored-by: Ken Wang <wangxiaotian2007@gmail.com> Co-authored-by: Spass <spass.ni42+github@gmail.com> Co-authored-by: zsun6 <115504839+zsun6@users.noreply.github.com> Co-authored-by: rishabh <rishabhsaiv@gmail.com> Co-authored-by: Yanzhao Wang <yzwang1996@gmail.com> Co-authored-by: devangpratap <115096812+devangpratap@users.noreply.github.com> Co-authored-by: Jasim Kareem <mj9034812@gmail.com> Co-authored-by: Ryan Choi <54298620+Ryan11c@users.noreply.github.com> Co-authored-by: Ronald Mannak <ronaldmannak@me.com> Co-authored-by: Mateusz Rajski <60450261+mateuuszzzzz@users.noreply.github.com> Co-authored-by: Ishaan Potle <ishaanpotle9@gmail.com> Co-authored-by: Moinuddin Shaik <monu.dev@proton.me> Co-authored-by: Moinuddin Shaik <moinuddinmoin1357@gmail.com> Co-authored-by: shaeantoine <42801108+shaeantoine@users.noreply.github.com> Co-authored-by: Nicoló Angileri <nicoloangileri635@gmail.com> Co-authored-by: Sam Small <sam.small@gmail.com> Co-authored-by: Ishti <45158028+ishtihoss@users.noreply.github.com> Co-authored-by: GokayAI <60583610+gokay-ai@users.noreply.github.com> Co-authored-by: Valeriy Sofin <112272443+sofinvalery@users.noreply.github.com> Co-authored-by: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Co-authored-by: Kaiyu Shi <skyisno.1@gmail.com> Co-authored-by: crabel99 <crabel@mac.com> Co-authored-by: nav <navneet.dagdiya@gmail.com> Co-authored-by: LongYinan <lynweklm@gmail.com> Co-authored-by: Joseph Heck <j_heck@apple.com> Co-authored-by: deadbydawn101 <sanjosepcrepair@gmail.com> Co-authored-by: deadbydawn <gothravenllm@gmail.com> Co-authored-by: Vijay Kapse <vijayskofficial@gmail.com> Co-authored-by: Long Yixing <longyixing331@gmail.com> Co-authored-by: Aditya Dhulipala <adi@dhulipala.xyz> Co-authored-by: Ishaan Mehta <106708534+ishaanzee@users.noreply.github.com> Co-authored-by: Mattt <mattt@me.com> Co-authored-by: Irakli Salia <65120973+Isalia20@users.noreply.github.com> Co-authored-by: Pasha Khosravi <khosravipasha@users.noreply.github.com> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> Co-authored-by: bri-prism <288398250+bri-prism@users.noreply.github.com> Co-authored-by: CodeAndCanvas728 <cedar-elk-01@icloud.com> Co-authored-by: SounLabs <159518213+SounLabs@users.noreply.github.com> Co-authored-by: David Dalcu <david.dalcu@gmail.com> Co-authored-by: Akshay Hitesh Shah <109842662+AkshayShah03@users.noreply.github.com> Co-authored-by: Sash <sash@ela.city> Co-authored-by: tc3oliver <57586652+tc3oliver@users.noreply.github.com> Co-authored-by: Magnus Lundstedt <magnus.lundstedt@infidyne.com> Co-authored-by: jrp2014 <jrp2014@users.noreply.github.com> Co-authored-by: Yi <39162855+SleepAviator@users.noreply.github.com> Co-authored-by: “Ronald <“Ronaldmannak@me.com”> Co-authored-by: Ivan Petrukha <93926277+petrukha-ivan@users.noreply.github.com> Co-authored-by: Haojin Yang <yanghaojin@hotmail.com> Co-authored-by: yanghaojin <haojin.yang@greenbit.ai> Co-authored-by: Marcin Markiewicz <75364832+jmmarkiewicz@users.noreply.github.com> Co-authored-by: Corentin Kerisit <corentin.kerisit@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Ambisphaeric <270742616+Ambisphaeric@users.noreply.github.com> Co-authored-by: Sankur Kundu <sankur.kundu.tw@gmail.com> Co-authored-by: begcdn <150344566+begcdn@users.noreply.github.com> Co-authored-by: Core <45273146+IamCooller@users.noreply.github.com> Co-authored-by: David Tivris <git@davidtivr.is>

zcbenz reviewed Sep 2, 2026

View reviewed changes

zcbenz mentioned this pull request Sep 3, 2026

Fix compile_dfs multi-output trace leaks #4440

Closed

protheus313 force-pushed the fix-multi-output-sibling-cycle branch from f034f8f to 78de35e Compare September 3, 2026 05:57

zcbenz reviewed Sep 3, 2026

View reviewed changes

Comment thread mlx/array.h Outdated

zcbenz approved these changes Sep 10, 2026

View reviewed changes

tudalex added 4 commits September 11, 2026 08:02

zcbenz force-pushed the fix-multi-output-sibling-cycle branch from dcc1057 to c489e6b Compare September 10, 2026 23:29

nit

e1259cd

zcbenz force-pushed the fix-multi-output-sibling-cycle branch from c489e6b to e1259cd Compare September 10, 2026 23:31

Fix crash in test

123cd15

zcbenz merged commit ae9df56 into ml-explore:main Sep 11, 2026
29 checks passed

aaishwarymishra pushed a commit to aaishwarymishra/mlx that referenced this pull request Sep 11, 2026

Break the sibling cycle when an array is released by assignment (ml-e…

3a7a2d4

…xplore#4453) Co-authored-by: Cheng <git@zcbenz.com>

zcbenz mentioned this pull request Sep 26, 2026

Memory leak when compiled function contains multi-output op #3932

Closed

This was referenced Sep 27, 2026

perf(glm5_next): cheaper exact GLM-5.3 decode steps (early eval, fused HC chain, routed gate/up) jundot/omlx#4019

Merged

Sync the final M5 Ultra stack: MiMo long-context attention, GLM compiled-FFN leak fix jonathan308/omlx-fusion#31

Merged

BrewTestBot mentioned this pull request Sep 29, 2026

mlx 0.32.3 Homebrew/homebrew-core#314148

Merged

2 tasks

jundot mentioned this pull request Oct 6, 2026

build: upgrade MLX to 0.32.3 with runtime compatibility jundot/omlx#4311

Merged

zcbenz mentioned this pull request Oct 19, 2026

Release the previous value through ~array() on assignment #4710

Closed

2 of 4 tasks

Navigation Menu

Repository navigation

Break the sibling cycle when an array is released by assignment - #4453

Break the sibling cycle when an array is released by assignment#4453
zcbenz merged 6 commits into
ml-explore:mainfrom
tudalex:fix-multi-output-sibling-cycle

tudalex commented Sep 2, 2026

Uh oh!

zcbenz left a comment

Uh oh!

zcbenz left a comment

Uh oh!

Uh oh!

tudalex commented Sep 3, 2026

Uh oh!

tudalex commented Sep 3, 2026

Uh oh!

zcbenz commented Sep 3, 2026

Uh oh!

zcbenz left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	// https://github.com/ml-explore/mlx/pull/1590
	TEST_CASE("test siblings circular references without eval") {
	std::weak_ptr<array::Data> tracker;
	auto fun = [&]() {
	array key({1, 2});
	auto splits = split(key, 2);
	{
	// Set fake data as a tracker for ArrayDesc's lifetime.
	splits[0].set_data(allocator::malloc(0));
	tracker = splits[0].data_shared_ptr();
	}
	auto a = reshape(splits[0], {});
	auto b = reshape(splits[1], {});
	return b;
	};
	fun();
	CHECK(tracker.expired());
	}

Conversation

tudalex commented Sep 2, 2026

Uh oh!

zcbenz left a comment

Choose a reason for hiding this comment

Uh oh!

zcbenz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tudalex commented Sep 3, 2026

Uh oh!

tudalex commented Sep 3, 2026

Uh oh!

zcbenz commented Sep 3, 2026

Uh oh!

zcbenz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants