fix(mlx): bump MLX pin to 81ba1c6a so mxfp8 block scales round up - #1772
Conversation
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
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
PR finalizationTest coverageAdded a regression test for the deferred drain's Failed branch: Documentation completenessChecked |
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
…#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.
Why
The fp8 round-trip bound failed on every Metal host, M1 Ultra byte-identically to M5 Max, because the pinned MLX
9a795735predates ml-explore/mlx#4353: Metal and CPU encoded the mxfp8 E8M0 block scale asround(log2(amax / 448)), so about half the blocks saturated their maxima, losing up to1 - 2^-1/2. CUDA rounds up, which is why #1742 passed on GB10. The test was right, andrequantize_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
9a795735to upstream main81ba1c6a(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.cppalso drops a leftoverelem_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.gather_qmmgainedglobal_scaleahead ofsorted_indices([Metal] global scale for qmm ml-explore/mlx#4458), so all 13 calls passstd::nullopt.link_cuda()names it,docs/installation.mdlists 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 insidefused_sampleor hides the error from its owner. A regression test aborts the process with the old drain and passes with the new one.fp8_block_requantize_round_trip_stays_within_half_an_e4m3_step, with the same seed and shape. It states the derivation and assertsgroup_max <= 448 * scaleper block.Validation (M1 Ultra, macOS 27.0)
OpenXLA feature linklinked a--features cuda,xla-ireerelease binary on GB10 at the new pin, which covers the overlays and the cuSOLVER link. The greenCUDA sm_70 compilecheck 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
array_evaluated_bytes, the server's lookahead read, is another non-Resultbridge function that now throws on a failed launch. It needs routing through the scheduler's step-failure path.metal/compiled.cppwould cast a comparison's inputs tobool. This is latent, since no compiled function contains a comparison.Closes #1769