Refresh mlx-swift fork and embedded MLX - #6
Merged
Conversation
* Add FP8 conversion utilities.
- add an iOS build in CI to catch iOS-only build breaks - move the docs build into the same macOS runner
* Add global scale support to quantized layers Store optional globalScale on QuantizedLinear and QuantizedEmbedding so nvfp4 weights can preserve the scale needed by lower-level MLX quantize/dequantize operations. Forward the scale when creating or dequantizing weights, add a direct pre-quantized QuantizedEmbedding initializer to match QuantizedLinear, and guard global-scale execution paths on Metal because MLX does not support globalScale dequantization there. Add focused tests for the new layer state and parameter exposure without broadening the generic quantization API surface. * Keep nvfp4 global scale on the quantized path for Metal Since MLX encodes nvfp4 group scales pre-multiplied by (448 * 6) / globalScale, the global scale factors out of dequantization linearly. Exploit this by always running quantizedMM (supported on Metal without globalScale) and applying globalScale / (448 * 6) to the output, instead of dequantizing the full weight and falling back to a plain matmul. This keeps weights quantized end-to-end on all backends, removes the dequantizedWeight fallback from QuantizedLinear and QuantizedEmbedding, and applies the scale before the bias add so the bias is unaffected. Add Metal tests that validate the GPU quantized path against CPU dequantized(globalScale:) references for QuantizedLinear (matmul + bias) and QuantizedEmbedding (asLinear and index lookup).
…er:) (ml-explore#448) The initializer retains a box holding the finalizer closure and hands it to mlx as an opaque payload, but `finalizerTrampoline` read it back with `takeUnretainedValue`, which never consumes the `passRetained` +1. The local `free(ptr:)` that would have balanced it was never passed anywhere — only the trampoline is — so it was dead code and the box leaked. The box owns the closure, so the closure's captures leak with it. The initializer's own documented example captures an `IOSurface` and comments that it "implicitly releases it when it returns"; it did not, and the surface was pinned for the lifetime of the process. `testIOSurface` covers the path but only prints, so nothing caught it. `takeRetainedValue` consumes the +1 exactly once. mlx-c documents the dtor as "Callback for when the buffer is no longer needed" and calls it once per buffer, and the trampoline has exactly one caller, so a single consume is correct. Removes the dead `free(ptr:)`. Adds a differential regression test: with the bug the finalizer runs but its captures are never released (witness releases 0); with the fix they are (1). The test drains mlx's allocator cache first, since a freed buffer's dtor is not guaranteed to have run at scope exit, and uses lock-guarded counters because mlx may invoke the dtor from its own thread. Full suite: 534 tests, 0 failures. Co-authored-by: Fred Jura <digital@afya.ai>
…them (ml-explore#445) `tools/update-mlx.sh` hand-tracked two lists that mlx's own CMake derives, so each one silently went stale whenever the vendored mlx submodule gained a jit source. The CUDA `MLX_JIT_SOURCES` argument named 16 `device/` headers explicitly, where `mlx/backend/cuda/CMakeLists.txt` uses `file(GLOB)` over `device/*.h` and `device/*.cuh`. A bump that adds one regenerated a `cuda_jit_sources.h` missing its array, failing later as `use of undeclared identifier 'jit_source_*'`: mlx v0.31.2 needs 17 and v0.32.0 needs 22, against the 16 the frozen list produced. The Metal step named 46 `make` targets explicitly, where `mlx/backend/metal/CMakeLists.txt` declares them via `make_jit_source()`. A bump that adds one left `jit/includes.h` declaring a `mlx::core::metal::<name>()` that no generated source defined, i.e. an undefined symbol at link time on macOS and iOS: v0.31.2 declares 47 and v0.32.0 declares 48, against the same 46. Both lists are now derived from the vendored tree at run time. Against the currently pinned mlx (v0.31.1) the derived values are identical to the removed literals, and the regenerated `cuda_jit_sources.h` is byte-for-byte identical to the committed one, so this is a no-op at the current pin. Each derivation is guarded: with `set -e` an empty result would not fail on its own, and bare `make` with no targets would quietly build `all` instead.
Advance the MLX submodule to the credited NVFP4 block-loader adaptation, regenerate both regular and NAX JIT kernel twins, and add a Metal equivalence test covering every e4m3 scale byte and packed fp4 byte (524,288 BF16 staged values).\n\nSource: mlx.fast submission 4228f9e1-18a5-4a1a-8047-96dc4e9289ca, validated commit b8fe3af.\n\nCo-authored-by: anupsv <6407789+anupsv@users.noreply.github.com>
a-github-name
marked this pull request as ready for review
August 14, 2026 13:51
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
da318704cc0e972b61dcca43c62cd15e545362aeb57bd7640f3f7c743b76a58478faaf1e8ee084f2Final revisions
3e6df6d8163a8f212061d15739eeeec12d5b89e3da318704cc0e972b61dcca43c62cd15e545362aeb57bd7640f3f7c743b76a58478faaf1e8ee084f2fb0c62d372d6aaa75edfbcb950d9dd797fce944a7df7bcde24dce2a672024be5Validation
swift-formatplusuvx pre-commit run --all-filesDependency
Linux/CUDA proof
The exact Swift and embedded-core revisions were built and executed on Ubuntu 24.04 arm64 with an NVIDIA GB10 (compute capability 12.1), driver 580.159.03, CUDA 13.3, Swift 6.0.3, cuDNN, NCCL, and OpenBLAS.
scripts/prepare-linux-native.shbuilt all 246 MLX/Swift bridge targets, andexample1 --device gpuran successfully onDevice(gpu, 0). The downstream full release CUDA package linked and launched with toolkit environment variables unset, resolved its CUDA/cuDNN/NCCL/NVRTC/cuFFT libraries, and selected native MLX GPU. The final mere.run Linux CPU, macOS Swift, app-bundle, docs, and security checks are green on the exact immutable pins.Downstream performance A/B
The dependent mere.run refresh was release-benchmarked against the previous production pin on an M4 Max. Laguna XS 2.1 remained flat at
107.76 -> 107.75 tok/sdecode with equal0.750 smedian prefill and byte-identical 128-token output. MiniMax-H3 remained flat at17.45 -> 17.51 smedian wall time, with its steady 50-block transformer evaluation0.828 -> 0.827 sand unchanged Metal memory. Full receipts are recorded on sawfwair/mere-run#297.