Merge upstream main and re-vendor the rebased core (numen-tech/mlx#10, gemma4-qat#282)

[jeethu]

Merges upstream ml-explore/mlx-swift main into prism and re-vendors the core from the rebased numen-tech/mlx fork (numen-tech/mlx#10: upstream ml-explore/mlx main 06eb7483f plus the fork's 10-patch series and review fixes). Tracking: numen-tech/gemma4-qat#282.

Commits

1. 084da487 Merge ml-explore/mlx-swift main 1e61d9c into prism. Upstream since the 0.32.2 merge-base 90194196: fix lock inversion between evalLock and dispatch_once on globalStreams ml-explore/mlx-swift#499 (evalLock / globalStreams lock inversion), fix xcframework builds ml-explore/mlx-swift#495, be able to trigger workflows on main. notes about upcoming releases ml-explore/mlx-swift#489, Fix attention mask for float16 ml-explore/mlx-swift#496 (float16 attention mask), fix MLXLogger ml-explore/mlx-swift#493 (MLXLogger, new cLogSupport target), package: exclude mlx/build cmake artifacts from Cmlx target ml-explore/mlx-swift#402 (exclude mlx/build), pick up fix-custom-io from mlx-c ml-explore/mlx-swift#476 (mlx-c fix-custom-io, mlx-c moves to ebc88f10). Two conflicts:
   • Source/MLX/Logging.swift: keeps the fork's osLogType(for:) (a warning maps to .default) and its metadata suffix, and takes upstream's is_os_log_enabled(osLog, level) check from fix MLXLogger ml-explore/mlx-swift#493.
   • Source/Cmlx/include-framework/Metal.hpp: a generated header whose conflict is only the metal-cpp commit stamp. Kept ours; the next commit regenerates it.
2. 8b127c7 bump numen-tech/mlx to 2fc35a34 (upstream main + fork patches, gemma4-qat#282); regenerate.
   • Source/Cmlx/mlx: d5da0546 (the 0.32.2 fork line) moves to 2fc35a34, the head of fork/282-sync-upstream-main. This includes the review fixes 87d00a604 (shorter NAX-gate comment in device.cpp) and 2fc35a344 (spec_decode_verify rejects non-integer draft_tokens, in mlx/fast.cpp).
   • Source/Cmlx/mlx-c: ebc88f10 moves to cfc471f4, from numen-tech/mlx-c. See the mlx-c section below.
   • mlx-generated/ and include-framework/ are regenerated with LC_ALL=C tools/update-mlx.sh.
   • MLX_VERSION is 0.32.4, from mlx/version.h at 06eb7483f. It is set in Package.swift and xcode/xcconfig/Cmlx.xcconfig.
   • Build plumbing for the files the core gained since 0.32.2 (details below).
   • An FP8 test fix (details below).
3. 031fd87 prism_readme: mlx-c comes from the numen-tech/mlx-c fork.
4. b721fdd Cmlx: build spec_decode into the metallib; fix the Linux CUDA sources for the new core. From review: adds spec_decode.metal to KERNEL_LIST (regenerated, two runs byte-identical); in the Linux CUDA branch, excludes the Windows-only cuda/wddm.cpp and links cusolver, as the core CMake does; refreshes the qmv_fast paragraph and version wording in prism_readme.md.

Fork patches kept

Patch	Where it lives now
1-bit and bias-free (affine_sym) exposure	In the core (patches 1–3 of numen-tech/mlx#10) and the regenerated mlx-generated/quantized*.cpp and metal/quantized*.h. No Swift-side code.
spec_decode_verify	Core patch 10. Its signature in mlx/fast.h is unchanged from d5da0546 through 2fc35a34; 2fc35a344 only adds input validation in mlx/fast.cpp. Here: the regenerated include-framework/mlx-fast.h declaration and, since b721fdd, the kernel in the SwiftPM metallib. No Swift binding.
MLX.GPUCounters and the mlx_counters_* extern "C" shim	Source/MLX/GPUCounters.swift and Source/Cmlx/mlx-counters/mlx_counters.cpp, both unchanged. mlx::core::metal::Counters / counters() / reset() in metal.h are identical at 2fc35a34.
Implied-bias QuantizedEmbedding gather	Source/MLXNN/Quantized.swift:234, unchanged.
FMT_CONSTEVAL	Already dropped on prism (6db5272: fmt 12.1.0). Nothing to keep.
global_scale	The original fix (PrismML mlx-c 699f32fe, which passed std::nullopt for unbound global_scale arguments) was subsumed by upstream mlx-c. The same class of fix is needed again for gather_qmm; it is now the mlx-c pin cfc471f4.

mlx-c decision

Pinned to numen-tech/mlx-c numen/0.7.0-gather-qmm-global-scale (cfc471f4 = v0.7.0 + ml-explore/mlx-c#137). Re-pin to upstream once ml-explore#137 is in a release.

Why a pin was needed: xcrun swift build with mlx-c at upstream's ebc88f10 fails.

Source/Cmlx/mlx-c/mlx/c/ops.cpp:1711:13: error: reference to type 'const std::optional<array>' could not bind to an lvalue of type 'bool'

• The cause is upstream mlx cbb4059d8 ("[Metal] global scale for qmm", #4458, in 0.32.3), which inserted global_scale before sorted_indices in gather_qmm.
• Moving to mlx-c v0.7.0 (a341b49) does not help: ebc88f10..v0.7.0 only bumps the version in CMakeLists.txt and docs/src/conf.py. No released mlx-c builds against MLX ≥ 0.32.3.
• Fix gather_qmm compatibility with MLX 0.32.3 ml-explore/mlx-c#137 is v0.7.0 plus a 4-line ops.cpp change. Under MLX_VERSION_NUMERIC >= 32003 it passes std::nullopt for the new argument. The C API is unchanged, so no Swift changes are needed.

Regeneration check

• Twice is byte-identical. Running LC_ALL=C tools/update-mlx.sh twice gives byte-identical mlx-generated/, include-framework/ and include/ (diff -r on snapshots: no output). This holds at both 78bdb20e and 2fc35a34.
• 78bdb20e and 2fc35a34 generate the same output. Neither device.cpp nor fast.cpp is a generator input, so mlx-generated/ has no diff between them.
• The output is independent of mlx-c. It is identical with mlx-c at ebc88f10 and at cfc471f4.
• One caveat, include-framework/Metal.hpp. metal-cpp's MakeSingleHeader.py stamps the superproject HEAD sha (Autogenerated from commit …). That line follows the commit the script runs on, not the submodule. The committed stamp is the merge 084da487. This is an upstream generator property; it is not patched.
• Every generated change traces to a core change. Every modified mlx-generated/metal/X has a changed mlx/backend/metal/kernels/X in d5da0546..2fc35a34. Every modified JIT .cpp embeds at least one header that changed in that range. compiled_preamble.cpp and cuda/cuda_jit_sources.h follow the core's CPU and CUDA device headers.
• Spot-check of quantized.cpp.
  • The quantized.h section of mlx-generated/quantized.cpp (after #line 1 "mlx/backend/metal/kernels/quantized.h") is identical to Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h.
  • mlx-generated/metal/quantized.h is identical to that header.
  • Both moved +154/−148 lines, the same as the header did.
  • Fork-symbol counts match between the generated file and the header: affine_sym 2/2, sym_derived_bias 3/3, qmv_wide 4/4, implied_bias 61/61, bits == 1 17/17.

Build plumbing for new core files

• tools/fix-metal-includes.sh (generator change). KERNEL_LIST mirrors the core's unconditional build_kernel() set, which now also has gated_delta_update, gated_delta_update_vjp, gather_mm_offsets and scaled_dot_product_attention_vjp. Without them those kernels are missing from the SwiftPM metallib at runtime. gather_mm dispatches gather_mm_offsets. b721fdd adds the fork's spec_decode as well. default.metallib now contains gather_mm_offsets, sdpa_vjp_*, gated_delta_update* and spec_decode_verify. fence.metal stays out (see Known gaps).
• Package.swift.
  • Excludes the new CUDA .cpp files (cusolver_utils, wddm) on Apple and Linux CPU builds; the macOS build fails without this (cuda_runtime.h not found). The Linux CUDA branch compiles cusolver_utils and links cusolver, and excludes the Windows-only wddm (b721fdd).
  • Excludes the new Metal .cpp files (cross_entropy, gated_delta_update, and the fork's spec_decode, which was previously missing) from the no-Metal Linux builds.
  • Drops the exclude for mlx/CODE_OF_CONDUCT.md, which the core deleted.
• xcode/MLX.xcodeproj. Adds membership exceptions for:
  • the new kernels/ and cuda/ files;
  • the new benchmark and python-test files and tests/cuda_tests.cpp;
  • the new generated metal/*.h headers;
  • the fork's doctest files tests/metal_counters_tests.cpp and tests/metal_kernel_tests.cpp. These already broke the xcodeproj build on prism 631cdf9e.
• Tests/MLXTests/FP8Tests.swift. Since Fix mx.from_fp8 E4M3FN NaN decode with branchless carry ml-explore/mlx#4376 (in the jump from 0.32.2 to 06eb7483f), 0x7F and 0xFF decode to NaN, which is correct for E4M3FN. The test's reference model expected ±480, so it now expects NaN and compares with equalNaN: true.

Verification (macOS, Xcode 27.0 / Swift 6.4)

Core at 2fc35a34, mlx-c at cfc471f4, head b721fdd (build and tests rerun after the review fixes, same counts):

• xcrun swift build: Build complete! The remaining Swift warnings are all in code identical to upstream or prism: withValue deprecations in Stream.swift/State.swift, plus MLXArray+maskFill.swift and MLXNN/Module.swift.
• xcodebuild test -scheme mlx-swift-Package -destination 'platform=macOS' -skipPackagePluginValidation -parallel-testing-enabled NO: ** TEST SUCCEEDED **.
  • MLXTests: 391 XCTest (2 skipped: distributed ring) and 8 Swift Testing.
  • MLXIntegrationTests: 669 Swift Testing.
  • CmlxTests: 1.
• The xcodeproj CI path: at core 78bdb20e, xcodebuild build-for-testing -scheme MLX -destination 'platform=macOS' reported ** TEST BUILD SUCCEEDED ** and xctest passed MLXTests (389 XCTest, 2 skipped; 8 Swift Testing) and MLXIntegrationTests (669). Rebuilt at b721fdd: ** TEST BUILD SUCCEEDED **. The 2fc35a34 delta adds no files.
• xcodebuild build -scheme mlx-swift-Package -destination 'generic/platform=iOS' at b721fdd: ** BUILD SUCCEEDED **.

Known gaps

• spec_decode_verify has no Swift or mlx-c binding; it never had one on prism. Since b721fdd its kernel is in the SwiftPM metallib, so a C++ caller in Cmlx works on GPU.
• fence.metal is not in the metallib, as in upstream mlx-swift. The core builds it only for Metal ≥ 3.2, and only MLX_METAL_FAST_SYNCH=1 (default off) uses it, so setting that variable fails with a missing kernel.
• Not verified:
  • the Linux CPU and Linux CUDA SwiftPM builds (no Linux host here; the CUDA changes mirror the core CMake);
  • the root CMakeLists.txt build, which still fetches upstream mlx v0.32.2 and mlx-c ebc88f10, as on prism. It is untouched here.

Codex round

• P1, spec_decode.metal missing from the metallib. GPU spec_decode_verify would fail with "Unable to load kernel". Fixed in b721fdd.
• P1, fence.metal missing from the metallib. Pre-existing, and the same in upstream mlx-swift. Kept as a known gap (above).
• Task review: the Linux CUDA branch compiled the Windows-only wddm.cpp and did not link cusolver. Fixed in b721fdd.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-11T14:54:23.706101Z	b721fdd	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b721fdd8f3

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add tvOS to cLogSupport's platform filters

When the Xcode project is built for tvOS, Logging.swift can import OSLog and therefore imports cLogSupport, but this framework link and the corresponding PBXTargetDependency both omit tvos from their platform filters. Since tools/create-xcframework.sh explicitly archives tvOS, the resulting tvOS build can lack the required Clang module/dependency; include tvos in both filters.

Useful? React with 👍 / 👎.