Skip to content

fix(mlx): bump MLX pin to 81ba1c6a so mxfp8 block scales round up - #1772

Merged
inureyes merged 6 commits into
mainfrom
fix/issue-1769-mlx-pin-mxfp8-round-up
Sep 10, 2026
Merged

inureyes merged 6 commits into
mainfrom
fix/issue-1769-mlx-pin-mxfp8-round-up

Conversation

@inureyes

@inureyes inureyes commented Sep 10, 2026

Copy link
Copy Markdown
Member

Why

The fp8 round-trip bound failed on every Metal host, M1 Ultra byte-identically to M5 Max, because the pinned MLX 9a795735 predates ml-explore/mlx#4353: Metal and CPU encoded the mxfp8 E8M0 block scale as round(log2(amax / 448)), so about half the blocks saturated their maxima, losing up to 1 - 2^-1/2. CUDA rounds up, which is why #1742 passed on GB10. The test was right, and requantize_block_fp8_weights, the only E8M0 quantize caller, was clipping vendor FP8 checkpoints on Metal. Widening the bound, as the issue proposed, was rejected.

What changed

  • MLX pin 9a795735 to upstream main 81ba1c6a (99 commits). The seven overlays whose targets upstream touched are three-way merged and keep their deltas; the other 21 are unchanged upstream. metal/compiled.cpp also drops a leftover elem_to_loc_1<uint> that undid part of Fix compiled kernel correctness for negative-strided inputs ml-explore/mlx#3720, so its only delta is the mixed-dtype cast.
  • Adaptations to what the new pin changes under the bridge:
    • gather_qmm gained global_scale ahead of sorted_indices ([Metal] global scale for qmm ml-explore/mlx#4458), so all 13 calls pass std::nullopt.
    • CUDA now needs cuSOLVER ([CUDA] Cholesky via cuSOLVER ml-explore/mlx#4208). link_cuda() names it, docs/installation.md lists it, and CI's link job now runs on pin and CUDA link-list changes, which is how this slipped past CI.
    • array::is_available() now throws and clears a failed launch's error ( Propagate CPU errors to events ml-explore/mlx#3742). The rejection sampler's drain now reads status, signal and error pointer instead, so a GPU fault no longer terminates the process from inside fused_sample or hides the error from its owner. A regression test aborts the process with the old drain and passes with the new one.
  • The round-trip check moves into its own test, fp8_block_requantize_round_trip_stays_within_half_an_e4m3_step, with the same seed and shape. It states the derivation and asserts group_max <= 448 * scale per block.

Validation (M1 Ultra, macOS 27.0)

  • fp8: the old pin fails the new test at block 0 (the maximum 4.8046875 scales to 615 and saturates). The new pin saturates 0 of 650 blocks, with a worst error of 0.0489 of the group max against 0.2928 before.
  • Turbo launchers pass (max RMS 1.7263e-4 and 1.5259e-4).
  • Teacher-forced logit traces, old pin vs new, on five checkpoints at widths 1, 8 and 256: 0 disagreements on decided positions. Over 4,096 positions on qwen3-30b-a3b, 1 of 1,720 decided positions differs, at the reference's rank 2, with perplexity -0.20%.
  • Branches the short runs missed:
    • head-dim-512 decode past 1,024 keys: identical text.
    • GQA-8 decode past 8,192 keys: identical text, decode 55.6 to 60.8 tok/s.
    • head-dim-72 vision towers: image prefill 239 to 279 tok/s. Descriptions diverge into equally faithful text; decided answers are unchanged.
  • Short-context throughput is within 0.6% on three checkpoints.
  • Workspace gate: 10986 passed, 0 failed, 359 ignored. Clippy, fmt, both pin parsers and the cross-repo reference check are clean.
  • CUDA was compiled and linked in CI but not run: OpenXLA feature link linked a --features cuda,xla-iree release binary on GB10 at the new pin, which covers the overlays and the cuSOLVER link. The green CUDA sm_70 compile check is a skip, because CUDA 13.0 cannot target sm_70.

The per-checkpoint numbers, method and derivations are in TECHNICAL_REPORTS/1772-mlx-pin-mxfp8-round-up-20260911.en.md.

Not validated, or reported and not fixed

  • No CUDA execution. No M5 Max run: the generation-17 NAX paths changed upstream in this range are unmeasured.
  • array_evaluated_bytes, the server's lookahead read, is another non-Result bridge function that now throws on a failed launch. It needs routing through the scheduler's step-failure path.
  • The mixed-dtype cast in metal/compiled.cpp would cast a comparison's inputs to bool. This is latent, since no compiled function contains a comparison.

Closes #1769

The fp8_block round-trip bound failed on every Metal host, M1 Ultra and M5 Max alike, because the pinned MLX (9a795735) predates ml-explore/mlx#4353. At that pin Metal and CPU encode the mxfp8 E8M0 block scale as round(log2(amax / 448)), which lands below amax / 448 for about half the blocks, so their maxima scale past 448 and saturate, losing up to 1 - 2^-1/2 of the block maximum. CUDA rounds up, which is why #1742 passed on GB10. The FP8 block requantize is the only E8M0 quantize caller, so vendor FP8 checkpoints on Metal were clipped the same way.

The pin moves to upstream main 81ba1c6a (99 commits). The seven overlays whose targets upstream touched in that range are three-way merged and keep their deltas; the other 21 are unchanged. ml-explore/mlx#4458 inserted global_scale ahead of sorted_indices in gather_qmm, so the 13 bridge calls now pass std::nullopt for it. The round-trip accuracy check moves out of the byte-identity test into its own test, which also asserts that no block maximum exceeds 448 times its scale.

Closes #1769
@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:high High priority area:models Model architectures, weights, loading, metadata area:core mlxcel-core: MLX FFI, primitives, KV cache, layers platform:macos macOS (Apple Silicon) specific status:review Under review labels Sep 10, 2026
Review of the pin bump found three things the new pin changes underneath this tree.

- ml-explore/mlx#4208 moved Cholesky onto cuSOLVER and `gpu::init()` now creates its handle cache on every CUDA start. MLX links it PRIVATE, so cargo never saw it and every `--features cuda` link would fail on `cusolverDnCreate`; `link_cuda()` now names `cusolver`.
- ml-explore/mlx#3742 made `array::is_available()` detach the event through `Event::check_error()`, which throws and clears a failed launch's error. The rejection sampler's deferred drain called it on slots other requests stashed, inside `fused_sample`, which is not a `Result` bridge, so a failed command buffer would terminate the process and hide the error from the request that owns it. The drain now reads status, signal and error pointer directly and drops a failed slot unread.
- The Metal `compiled.cpp` overlay still emitted `elem_to_loc_1<uint>` for 1-D inputs, half of ml-explore/mlx#3720 that an earlier sync missed; it now matches upstream, so the overlay's only delta is the mixed-dtype cast. The CUDA mixed-type `FloorDivide` overload floors like upstream's float branch (ml-explore/mlx#4108), and three stale sync notes are corrected.

Workspace gate 10985 passed, 0 failed; clippy and fmt clean on Metal. The CUDA link is not verifiable on this host.

Refs #1769
The MLX pin bump in this branch makes the CUDA backend link and initialize cuSOLVER, so a prebuilt CUDA binary now needs libcusolver on the host. The runtime requirements section listed only headers; it now names the shared libraries the binary links, cuSOLVER included, and the loader error a runtime-only install produces without it.

Refs #1769
drain_pending_verification's Failed path (stashed_launch_state reading status, the event signal and its error pointer instead of array::is_available()) had no regression test, only the reasoning in its comment. Reverting it to call is_available() on the fixture aborts the whole test process with an uncaught std::runtime_error, exactly the ml-explore/mlx#3742 failure mode this PR works around, so a silent reintroduction would not fail by name without this test.

Adds a test-only cxx bridge hook (sampling_dispatch_stash_failed_launch_for_test, plus two accessors) that builds a valid, signalled event carrying an error through MLX's public Event/Error API, stashes it into pending-verification slot 0, then asserts the drain neither throws nor consumes the error and drops the slot. Confirmed by temporarily reverting drain_pending_verification to is_available() and watching the new test crash the process; restored afterward.

Validation:
- cargo test --profile test-fast --features metal,accelerate -p mlxcel-core --lib sampling_rejection_tests:: (28 passed)
- cargo fmt --check and cargo clippy -p mlxcel-core --lib -- -D warnings (clean, narrow scope per CLAUDE.md)

Refs #1769
@inureyes

Copy link
Copy Markdown
Member Author

PR finalization

Test coverage

Added a regression test for the deferred drain's Failed branch: a_failed_stashed_launch_is_dropped_without_throwing_or_consuming_its_error in src/lib/mlxcel-core/src/sampling_rejection_tests.rs, backed by a new test-only cxx bridge hook (sampling_dispatch_stash_failed_launch_for_test plus two accessors) in mlx_cxx_bridge.cpp/.h and lib.rs, following the existing sampling_dispatch_drain_pending/sampling_dispatch_reset hook pattern. It builds a valid, signalled event carrying an error through MLX's public Event/Error API (no GPU JIT needed), stashes it into pending-verification slot 0, then asserts drain_pending_verification neither throws nor consumes the error and drops the slot. Proved by temporarily reverting the drain to call is_available(): the reverted build crashed the whole test process with libc++abi: terminating due to uncaught exception of type std::runtime_error: mlxcel test: synthetic launch failure, the exact ml-explore/mlx#3742 failure mode this PR works around, then restored. cargo test --profile test-fast --features metal,accelerate -p mlxcel-core --lib sampling_rejection_tests:: reports 28 passed with the fix restored. Commit a2db793.

Documentation completeness

Checked docs/supported-models.md (Qwen3.5/3.8 FP8 entry), docs/architecture.md, CONTRIBUTING.md, README.md, docs/environment-variables.md, docs/turbo-kv-cache.md, and docs/upstream/mlx-cuda-quantized-matmul-nondeterminism.md for stale MLX pin, CUDA link library, or fp8-accuracy claims. None name the old pin as a current-state fact or make an accuracy claim the pin bump invalidates; mentions are either mechanism descriptions that hold regardless of the pinned commit, or historical citations (CHANGELOG.md's 9a795735 entry, CONTRIBUTING.md's 2c46b953 CUDA-threading citation, the upstream bug report pinned to b7c3dd6d) that are provenance and correctly left alone, matching how CHANGELOG.md is only touched at release-prep commits rather than per-PR. docs/installation.md remains the one doc that needed the cuSOLVER update, already in place. No changes made beyond the test coverage commit.

The xla-link job is the only PR job that links a `--features cuda` binary, and its path filter covered the IREE half of the link line but not the CUDA half. `src/lib/mlxcel-core/build.rs` names the CUDA libraries, and the MLX pin in `src/lib/mlx-cpp/CMakeLists.txt` decides which ones `libmlx.a` needs. This branch's pin bump added cuSOLVER (ml-explore/mlx#4208) without touching any build script, and the only CUDA job that ran was `cargo check`, which never links, so the missing library was found by review rather than CI. Both paths now trigger the job, which also makes it verify this branch's cuSOLVER link on GB10.

Refs #1769
@inureyes inureyes added status:done Completed and removed status:review Under review labels Sep 10, 2026
@inureyes
inureyes merged commit 8ad62d4 into main Sep 10, 2026
13 checks passed
@inureyes
inureyes deleted the fix/issue-1769-mlx-pin-mxfp8-round-up branch September 10, 2026 18:44
inureyes added a commit that referenced this pull request Sep 10, 2026
…#1778)

## Summary

The PR #1753 technical report listed three defects in `split-mtp` among its follow-ups. This PR fixes all three. Review of the first version found five more gaps in the same guards, and those are fixed here as well.

- **Refusal text.** The `MAX_BLOCK_SIZE` refusal rendered with two runs of 14 literal spaces. It is now one sentence with single spaces.
- **`--force` never deletes weight files.** `prepare_output_dir(dir, force)` refuses and names the file when the output directory holds any `*.safetensors` other than the drafter's own `model.safetensors`. The loader's `glob_safetensors` reads every `*.safetensors` in a directory, so a stale `model-0000N-of-...` or `consolidated.safetensors` would otherwise be loaded as part of the drafter. A stale `model.safetensors.index.json` is still cleared when no foreign file remains. An unreadable directory fails closed.
- **Quantization parameters are checked before any tensor work.** `SUPPORTED_AFFINE_BITS` (2, 3, 4, 5, 6, 8) and `SUPPORTED_AFFINE_GROUP_SIZES` (32, 64, 128) now live in `mlxcel-core::layers`, next to `validate_affine_quantization_bits` and `validate_affine_quantization_group_size`. Before this PR, `--q-bits 1` or `--q-group-size 16` did the tensor work and then aborted the process inside MLX's quantize. Gemma 4's private copy of the bit set now reads the shared constant. `validate_quantization_params` still accepts `1..=32` at load time (#929).
- **Checks run before mutation.** `preflight_checks` (bits, group size, output equals source) runs before `prepare_output_dir`. A bad flag or `-o` pointing at the source can no longer touch the directory first.

## Validation

- Workspace gate on 059ea43: exit 0, 123 binaries, 11,038 passed, 0 failed. Clippy, fmt, `verify-versions`, `verify-kernel-dtype-keys` and `verify-llama-compat` are clean. CI is green.
- Every new or strengthened test was run with its fix reverted and failed by name. With only the load-time `1..=32` check left, `--q-bits 1` aborts the test binary (`std::invalid_argument: [quantize] The requested number of bits 1 is not supported`, SIGABRT).
- Real run with a release binary at 059ea43 against `glm-4.7-flash-bf16`:
  - `--q-bits 4` writes the drafter: 54 tensors, 0.72 GB, `glm4_moe_lite_mtp`, 4-bit affine, group size 64. `model.safetensors` and `config.json` are byte-identical to the output of the pre-PR binary (the #1772 merge).
  - `--force` over a planted `model-00001-of-00003.safetensors` exits 1 and names that file. Every file's name, size, mtime and inode is unchanged.
  - `--q-bits 7` exits 1 within a second with `quantization bits (7) must be one of 2, 3, 4, 5, 6, 8`, and no `drafter2` is created.

Closes #1763.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:models Model architectures, weights, loading, metadata platform:macos macOS (Apple Silicon) specific priority:high High priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

test(fp8_block): fp8_block_requantize_matches_direct_path fails deterministically on M5 Max with macOS 27 / Xcode 27

1 participant