Skip to content

[CUDA] Add global scale support to gather_qmm - #4507

Merged
nastya236 merged 2 commits into
ml-explore:mainfrom
dhiltgen:cuda_g_qmm_global_scle
Sep 17, 2026
Merged

[CUDA] Add global scale support to gather_qmm#4507
nastya236 merged 2 commits into
ml-explore:mainfrom
dhiltgen:cuda_g_qmm_global_scle

Conversation

@dhiltgen

@dhiltgen dhiltgen commented Sep 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

This adds support for global scales in gather_qmm on cuda.

mlx-lm perf data with qwen 3.6 35B tested on a 5090. (p2048/g128)

model / system prompt tps gen tps
mlx-community/Qwen3.6-35B-A3B-nvfp4 (no global scale) 394.3 75.9
Model Optimizer export (with global scale) 404.8 70.5
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: co-developed with a coding agent

Attached wrapper script loads the Model Optimizer exports and wires up gather_qmm global scale, and adjusts dtypes for a fair comparison with the standard community model.
modelopt_nvfp4_mlxlm.py

@dhiltgen dhiltgen mentioned this pull request Sep 14, 2026
Merged
@nastya236

nastya236 commented Sep 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Thanks! Why "optimizes performance"? Did you notice a speed up?

General overall note: after giving this some more thought, I realised that dequantization during inference would be correct if and only if it matches the training. I believe there are two possible ways to produce nvfp4 weights:

  1. quantize(W) -> use nvfp4 x nvfp4 instructions for matmul and fold global scale into epilogue — in this case we also must fold scale into epilogue during inference as well (like in this PR)
  2. If the global scale was applied when dequantizing the weights, before casting to BF16 and performing the matrix multiplication in bf16 (quantization aware training on sm < 100), we should follow that same sequence during inference — this is done in [Metal] global scale for qmm #4458

Given that in the future all quantized operations during training will be native, I think it is okay to agree on folding it into epilogue. However, probably we should do the same for metal..

@nastya236
nastya236 self-requested a review September 16, 2026 09:56
nastya236
nastya236 requested changes Sep 16, 2026
View reviewed changes

@nastya236 nastya236 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some tests are failing can you have a look?

@nastya236 nastya236 added the await response This pull request is waiting for response from the author. label Sep 16, 2026
@dhiltgen
dhiltgen force-pushed the cuda_g_qmm_global_scle branch from fcd135f to d3012f2 Compare September 16, 2026 16:49
@dhiltgen

dhiltgen commented Sep 16, 2026

Copy link
Copy Markdown
Contributor Author

Why "optimizes performance"? Did you notice a speed up?

On further testing, I couldn't reproduce any speedup in the second commit so I've dropped it and the claim.

@nastya236
nastya236 merged commit 59d600b into ml-explore:main Sep 17, 2026
29 checks passed
@dhiltgen
dhiltgen deleted the cuda_g_qmm_global_scle branch September 17, 2026 15:06
@rcfa rcfa mentioned this pull request Oct 13, 2026
Open
@inureyes inureyes mentioned this pull request Oct 17, 2026
Merged
inureyes added a commit to lablup/mlxcel that referenced this pull request Oct 17, 2026
@inureyes
fix(moe): load mlx-lm NVFP4 per-expert global_scale natively (#2351) (#…
1a22cb0 
…2441)

## Summary

mlx-lm's `quantize_model(..., mode="nvfp4", global_scale=True)` (ml-explore/mlx-lm#1844) writes `max|w[e]|` per expert as `{switch layer}.global_scale` and flags the layer with `"global_scale": true`; MLX's `gather_qmm` dequantizes with `global_scale / (448 * 6)`. mlxcel read the same key as the compressed-tensors output multiplier, so such an export would decode with every routed expert 2688 times too large. This PR keeps the two scales apart, threads the native one through `gather_qmm`, and extends ModelOpt NVFP4 expert loading beyond Gemma 4.

### Part 1: MLX-native per-expert global scale

- FFI: `gather_qmm` gains a nullable `global_scale: *const MlxArray` (Rust decl, `mlx_cxx_bridge.h`, `mlx_cxx_bridge.cpp`); every existing caller passes null. A test-only `quantize_weights_with_global_scale` wraps `mlx::core::quantize(..., global_scale)`.
- `SwitchLinear::Quantized` now has `inv_global_scale` (the post-matmul multiplier, from the in-process `{prefix}.inv_global_scale` key) and `native_global_scale` (handed to `gather_qmm`). `quantized_parts` returns `None` when either is present, so the fused paths fall back to `forward`; `SwitchGLU::forward_with_expert_scales` passes through `forward_indexed` and gets the native scale too.
- The flag comes from the config, not the key. `apply_native_global_scale_flags` (called by `load_text_weights` and by every Gemma 4 loader through `sanitize_gemma4_nvfp4_weights`) reads the `quantization` / `quantization_config` entries (top level and `text_config`) and moves each flagged `{path}.global_scale` to `{path}.native_global_scale`. A flag without its tensor fails there, and so does a flag on a layer whose weight is not a `[num_experts, out, in]` stack (its scale would otherwise be renamed and never read); a `.global_scale` left at a switch prefix (never flagged) fails in `SwitchLinear` with the prefix named; a checkpoint that ships `native_global_scale` itself is refused. I chose this over adding a parameter to `from_weights_with_mode` because about 55 family files construct `SwitchLinear` / `SwitchGLU`; the decision is still made from the per-layer config entry, once, before any family sanitizer renames keys. Loaders that do not go through these funnels (other VLM paths) see the bare key; it fails the load wherever it sits at a shared `SwitchLinear` prefix, but a loader that renames expert keys by suffix can leave it orphaned (see Not done here).
- Load-time validation of the native scale against MLX's preconditions: `nvfp4` mode, one f32 entry per expert of a `[E, out, in]` stack, a Metal or CUDA GPU stream. Each failure is a load error naming the prefix and backend, never a mid-generation throw.
- Laguna's in-process expert sidecar is renamed to `inv_global_scale` (`laguna_sanitize.rs`, `laguna.rs` detection, tests). Laguna's dense shared-expert `[1]` `.global_scale` keeps its name: it feeds `UnifiedLinear`, and dense global scales are out of scope.
- Two fixes needed for any mlx-lm nvfp4 MoE export to load at all: `SwitchLinear::from_weights` (the shared loader, which declares no mode) infers bias-free block-float planes with uint8 scales from the shapes (nvfp4 / mxfp4 / mxfp8 by words-per-group) instead of declaring them affine and refusing them; and the quantized embedding lookup in the bridge stops requesting a uint8 output dtype for block-float tables, which aborted the first lookup (`[dequantize] Only real floating types are supported but output dtype == uint8`).

### Part 2: ModelOpt NVFP4 experts outside Gemma 4

- `load_text_weights` runs `normalize_nvfp4_keys` + `repack_nvfp4_weights_to_quantized` for any text checkpoint declaring ModelOpt NVFP4 in `quantization_config` (top level or `text_config`) or in `hf_quant_config.json`. Detection reuses Gemma 4's `is_modelopt_nvfp4`, lifted to a module-level `pub(crate)` fn; the `hf_quant_config.json` shape is rewritten into the `config.json` shape before it is checked. `quant_method` and the producer name match case-insensitively.
- The repack decodes F8_E4M3 block scales that MLX's own safetensors loader returns as raw uint8 bytes (the Gemma 4 loader already decoded them; a numeric cast would read the bytes as 0..255).
- New shared helper `stack_expert_leaf` plus `EXPERT_PLANE_LEAVES` in `switch_layers.rs`; the DeepSeek-V3, Mistral 4 and PhiMoE stacking loops and the generic per-expert stacker use it, so the repack's per-expert `weight_scale_2` sidecar becomes a `[num_experts]` `inv_global_scale` vector in the same pass as weight/scales/biases.
- PhiMoE and Mistral 4 read the 16-wide block size from `quantization_config` (`modelopt_nvfp4_group_size`). DeepSeek-V3's own expert path is affine-only; it now refuses a block-float plane by name instead of reporting a missing `.biases` key.
- Families with their own `weight` / `scales` / `biases` stacking loop that load through the shared `SwitchLinear` (OLMoE, dots1, Solar Open, LFM2, Cohere2-MoE, ...) would stack a ModelOpt plane's codes and block scales and drop its `weight_scale_2`. `SwitchLinear` now refuses a bias-free uint8-scale plane with no multiplier while per-expert `.global_scale` sidecars remain under its MoE block, naming the leftover key, so those checkpoints fail to load instead of decoding wrong.
- gpt-oss's own `ExpertLinear` calls `gather_qmm` with no tensor scale, so it refuses a plane carrying `.native_global_scale` or a bare `.global_scale`.
- `docs/supported-models.md` documents both formats.

Also added: `scripts/mlx_lm_logit_trace.py`, an mlx-lm reference tracer in the `examples/logit_trace` TSV format, reusing the tokenizer and chunk plan from `scripts/laguna_oracle_trace.py`.

## Verification

All on the GB10 CUDA host (sm_121), `--profile test-fast --features cuda`. MLX at the pin supports `gather_qmm(..., global_scale)` for nvfp4 on CUDA (overlay `quantized.cpp`, ml-explore/mlx#4507 plumbing; qmv and qmm_naive apply it, qmm_sm80 and the grouped-GEMM path decline it), confirmed first with a Python probe on mlx 0.32.4 and then by the Rust tests below. Metal is left to the Metal nightly. Checkpoint paths are under `~/models_archive/mlx` because `~/models` was remapped on the host during the run.

- New tests, all run on CUDA: `switch_layers::global_scale_tests` (native path against `x @ dequantize(w[e]).T * gs/2688`, unsorted and sorted layouts, plus a relative-L2 check against the float weights; flag stripped fails to load; flag without tensor fails; reserved key refused; nvfp4 / shape / exclusivity checks; ModelOpt sidecar stacking and application; partial sidecar refused), `sanitize_tests::load_text_weights_repacks_modelopt_nvfp4_for_non_gemma_families` (config.json and hf_quant_config.json detection, E4M3 bytes preserved), `load_text_weights_resolves_flagged_native_global_scales`, `load_text_weights_leaves_undeclared_triplets_alone`, and per-family stacking tests in `phimoe`, `mistral4`, `deepseek_v3`.
- Revert checks. With part 1 reverted (native scale applied as the old post-multiplier, bare key accepted, flags ignored): `native_global_scale_matches_dense_dequantized_reference`, `global_scale_without_the_config_flag_fails_to_load`, `config_flag_without_global_scale_tensor_is_a_load_error`, `native_global_scale_needs_nvfp4_and_one_entry_per_expert` and `load_text_weights_resolves_flagged_native_global_scales` fail. Without the gpt-oss guard, `gpt_oss_expert_linear_refuses_a_per_expert_global_scale` fails; without the unstacked-sidecar guard, `a_plane_whose_family_left_the_modelopt_sidecars_unstacked_is_refused` fails. With part 2 reverted (no sidecar leaf, no F8 decode): the PhiMoE, Mistral 4, DeepSeek-V3 stacking tests, `per_expert_sidecars_stack_into_the_output_multiplier`, `a_partial_per_expert_sidecar_is_refused` and `load_text_weights_repacks_modelopt_nvfp4_for_non_gemma_families` fail.
- Module suites (`--test-threads=1`): switch_layers, laguna, sanitize, phimoe, mistral4, deepseek_v3, deepseek, gemma4, olmoe, qwen3_moe, inkling, kimi_k3, deepseek_v4, gpt_oss, glm4_moe, step3p5, nemotron_h, hunyuan, llama4, exaone, ernie, qwen3_next all pass. `cargo clippy --profile test-fast --features cuda --lib --tests -- -D warnings`, `cargo check -p mlxcel-core --lib --tests`, `cargo fmt --check` and the five contract checks pass.

### Real checkpoints

- Laguna regression, `~/models_archive/mlx/laguna-xs-2.1-nvfp4` (`poolside/Laguna-XS-2.1-NVFP4`), now on `inv_global_scale`: `MLXCEL_SDPA_DETERMINISTIC=1 mlxcel generate -m <dir> -p "Explain how mixture-of-experts routing works in two paragraphs." -n 64 --temp 0 --show-reasoning` is token-identical to `main` (8744c87, built in the same worktree before the change; main's self-parity checked first).
- ModelOpt path that already worked, `~/models_archive/mlx/gemma-4-31b-it-nvfp4`: same command, token-identical to `main`.
- Real mlx-lm export, `~/models_archive/mlx/olmoe-1b-7b-0924-nvfp4-gs`. Recipe: `allenai/OLMoE-1B-7B-0924` (HF bf16; its config lacks `rms_norm_eps`, so `1e-5`, the transformers default, was added for mlx-lm), loaded with `mlx_lm.utils.load(..., lazy=True)`, floats cast to bf16, `quantize_model(model, config, None, None, mode="nvfp4", global_scale=True)`, `save(...)`. mlx-lm 0.32.2 (mlx 0.32.4, `mlx-cuda-13` wheel) already carries the `global_scale` path, so it was used instead of a `ba8a08f` checkout; 48 switch layers carry the flag. `main` refuses it at load (`quantization.mode is "affine" but the checkpoint ships no .biases plane`). This branch: `mlxcel generate -m <dir> -p "The capital of France is Paris. The capital of Germany is" -n 40 --temp 0` gives "Berlin. The capital of the United Kingdom is London. The capital of the United States is Washington, D.C." (mlx-lm greedy: "Berlin. The capital of Italy is Rome. ..."). With the config flags stripped it fails to load naming `model.layers.0.mlp.switch_mlp.gate_proj`.
- Teacher-forced logit trace on that export, corpus `docs/architecture.md`, `256 4 8 0` (width 256, 4 chunks, top-8), reference `scripts/mlx_lm_logit_trace.py` (mlx-lm 0.32.2 on CUDA), compared with `scripts/compare_logit_traces.py`:
  - fixed arm: 0 of 396 decided positions disagree (gap >= 2), 19 of 1024 overall, all at reference gaps under 0.5; perplexity 20.15 vs 20.13.
  - native scale replaced by the old post-multiply: 396 of 396 decided positions disagree (1024 of 1024 overall); greedy output is garbage ("ByrPLAYpackavelorschKK Television...").
- ModelOpt non-Gemma: no ModelOpt NVFP4 MoE checkpoint of DeepSeek-V3, Mistral 4 or PhiMoE is published at a size this host fits (`mistralai/Mistral-Small-4-119B-2603-NVFP4` is in Mistral's consolidated format), so the synthetic triplet tests above cover the stacking. I also tried `nvidia/Qwen3-30B-A3B-NVFP4` (`~/models_archive/hf-src/Qwen3-30B-A3B-NVFP4`): the 18624 triplets now repack and the experts stack, but the model still aborts in the dense fused QKV projection (`FusedQKVLinear` hardcodes affine and has no global-scale path), which is outside this issue; `main` aborts on the same checkpoint with a matmul shape error.

## Not done here

- Dense fused QKV with NVFP4 + `weight_scale_2` (blocks Qwen3-MoE ModelOpt checkpoints).
- VLM loaders other than Gemma 4 do not resolve the flag. An mlx-vlm export with per-expert global scales (Blaizzy/mlx-vlm#2437) fails to load where the bare key reaches a shared `SwitchLinear` prefix, but the Gemma 4 unified loader's `split_fused_moe_key` (`src/loading/vlm_gemma_unified.rs`) splits only `.weight` / `.scales` / `.biases` of a fused `experts.gate_up_proj` and would leave a `.global_scale` orphaned. That is pre-existing and in a loading file outside this issue.
- A security review comment on this PR lists LOW follow-ups (load-time key scan cost, flag resolution only in the text and Gemma 4 funnels, DeepSeek v1's private loader treating a bias-free quantized plane as unquantized, which predates this PR).
- Re-measured after review: Laguna XS 2.1 greedy is still token-identical to `main`, and the OLMoE export's greedy output is unchanged.

Closes #2351
@georgeguimaraes georgeguimaraes mentioned this pull request Oct 18, 2026
Merged
@Sofille65 Sofille65 mentioned this pull request Oct 20, 2026
Merged
Gajesh2007 pushed a commit to Layr-Labs/mlx that referenced this pull request Oct 23, 2026
@davidtai @kapellirohith
@zcbenz @ternaus @deBrian07 @erwinzhang7 @dhiltgen Daniel Hiltgen @aaishwarymishra @strayberry @dwijenpatel @nastya236 @ayaangazali @katlun-lgtm katlun-lgtm @Adityaj0 @eyupcanakman @prady0t @jagrit06 @PhilipJohnBasile @aleroot @dudududukim @atirna @saud5150 Saud @roshaninfordham @sashko-zakharchuk @XXXXRT666 @vaibhav8a @TheDarkchip @2sumtech @Ved235 @JasonHonKL @simeetnayan81 @Cdimoy @RohanGautam @michalk8 @nileshpatil6 @MarcosAsh @pierre427 Pierre Lamy @wyanzhao @davidkoski @axiom-of-choice @tudalex @Ritabanm @WindChimeRan @tpegolotti @Drifter4242 @Drifter42 @louen @keeeeenw @sp4s-s @zsun6 @rishabhsai @devangpratap @Prudctual @Ryan11c @ronaldmannak @mateuuszzzzz @IshaanPotle @CodeWithMoin @shaeantoine @nicoloangileri @ssmall256 @ishtihoss @gokay-ai @sofinvalery @devYRPauli @Stonesjtu @crabel99 @navthings @Brooooooklyn @heckj @DeadByDawn101 deadbydawn @vijay-kapse @Lyxot @adi-dhulipala @ishaanzee @mattt @Isalia20 @khosravipasha @angeloskath @bri-prism @CodeAndCanvas728 @SounLabs @ddalcu @AkshayShah03 @SashaMIT @tc3oliver @magnuslundstedt @jrp2014 @SleepAviator “Ronald @petrukha-ivan @yanghaojin yanghaojin @jmmarkiewicz @cerisier @claude @Ambisphaeric @Shakhtar-Sankur @begcdn @IamCooller @tivris
Merge upstream ml-explore/mlx main afc03c5 (2026-10-23) (#38)
7b0ebaf 
* 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@nastya236 nastya236 nastya236 approved these changes

Assignees

No one assigned

Labels

await response This pull request is waiting for response from the author.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@dhiltgen @nastya236