chore: bump MLX pin to 0.32.1 (#703) by inureyes · Pull Request #704 · lablup/mlxcel

inureyes · 2026-07-09T02:17:15Z

Summary

Bump all three synchronized MLX pin locations to upstream 57c66cac7cb3e5b1eb350488a61f1506b40d39f8 (Patch bump to 0.32.1, Patch bump to 0.32.1 ml-explore/mlx#3816).
Fix mlxcel's Metal compiled patch overlay for the latest compiled_collapse_contiguous_dims return shape by binding negative_strides, selecting the large-index path for negative strides, and generating the rank-1 large strided Metal kernel that upstream MLX now needs.
Note the relevant upstream range includes Fix fp quantized matvec for output dim < 8 ml-explore/mlx#3804 (Fix fp quantized matvec for output dim < 8), Add math mode option for custom Metal kernels ml-explore/mlx#3728 (Add math mode option for custom Metal kernels), and Fix compiled kernel correctness for negative-strided inputs ml-explore/mlx#3720 (Fix compiled kernel correctness for negative-strided inputs).

Test plan

cargo build --release --bin mlxcel-bench-decode
cargo test --release -p mlxcel-core sparse_v_kernel_threshold_zero_matches_graph
cargo test --release -p mlxcel-core delegated_fused_kernel_matches_reference_over_200_steps
cargo test --release -p mlxcel-core delegated_steel_envelope_matches_cold_only_fused_over_200_steps
python3 scripts/ci/check_cross_repo_refs.py

Closes #703

Update the three synchronized MLX pin locations to upstream commit 57c66cac7cb3e5b1eb350488a61f1506b40d39f8 and repair the Metal compiled overlay for the latest compiled_collapse_contiguous_dims return shape. The overlay now matches upstream CPU/CUDA by binding negative_strides and forcing the large-index path when negative strides are present. Validation: - cargo build --release --bin mlxcel-bench-decode - cargo test --release -p mlxcel-core sparse_v_kernel_threshold_zero_matches_graph - cargo test --release -p mlxcel-core delegated_fused_kernel_matches_reference_over_200_steps - cargo test --release -p mlxcel-core delegated_steel_envelope_matches_cold_only_fused_over_200_steps - python3 scripts/ci/check_cross_repo_refs.py Closes #703

inureyes · 2026-07-09T02:19:48Z

inureyes · 2026-07-09T02:20:05Z

Blocking regression risk found in the Metal compiled overlay.

compiled.cpp now sets large = true when negative_strides is present, so a non-contiguous 1-D negative-strided compiled op will request <kernel_lib>_strided_1_large. However the overlay still only builds _strided_${i}_large when i > 1, so _strided_1_large is missing. Upstream MLX at 57c66cac7cb3e5b1eb350488a61f1506b40d39f8 removed that guard and explicitly generates the large-index variant for every ndim including ndim=1.

Please mirror upstream here by generating _strided_1_large as well (remove the if (i > 1) guard / keep the upstream comment). The three MLX pin locations otherwise match, and I did not find local path leakage or a release workflow pin mismatch in the changed files.

Generate the large strided compiled kernels for every static rank in the local Metal overlay so negative-strided rank-1 inputs can select the large-index path that upstream MLX now exposes through compiled_collapse_contiguous_dims. Validation: - cargo build --release --bin mlxcel-bench-decode - cargo test --release -p mlxcel-core sparse_v_kernel_threshold_zero_matches_graph - cargo test --release -p mlxcel-core delegated_fused_kernel_matches_reference_over_200_steps - cargo test --release -p mlxcel-core delegated_steel_envelope_matches_cold_only_fused_over_200_steps - python3 scripts/ci/check_cross_repo_refs.py Refs #703

inureyes · 2026-07-09T02:31:07Z

Resolved in 80e7bd548: the local Metal overlay now generates _strided_1_large by removing the i > 1 guard, matching upstream MLX behavior for negative-strided rank-1 compiled inputs. Re-ran the AGENTS.md fused Metal tests plus cargo build --release --bin mlxcel-bench-decode successfully.

…711) ## Summary Complete the backend-aware fused-MoE `Dff` cap for #330 by recording the honest provenance of the CUDA default and pinning the decline boundary with a test. The cap-resolution logic itself (Metal 4096 / CUDA 8192, the `MLXCEL_FUSED_MOE_MAX_DFF` override, the pure `fused_moe_max_dff_from`, and the 2026-07-03 GB10 sweep) already landed in #643; that work was measured under MLX pin `e9463bb` (#626) and was never linked to #330. Since #643 the MLX pin advanced to 0.32.1 (#703/#704, commit `57c66cac`), so the CUDA default is now a carried-forward documented-data value, not one confirmed on current binaries. This PR closes that honesty gap, documents the runtime-vs-`cfg!` backend choice, adds a decline-boundary test, and formally closes #330. The cap value 8192 is unchanged. ## What changed - `src/models/switch_layers.rs`: expanded the `FUSED_MOE_MAX_DFF_CUDA` doc comment with a provenance caveat (the 8192 sweep ran on MLX pin `e9463bb`/#626; the pin has since moved to 0.32.1/#703/#704, commit `57c66cac`; 8192 is the conservative crossover floor kept so pin drift cannot flip it into a regression; a GB10 re-validation on 0.32.1 is pending). - `src/models/switch_layers.rs`: documented on `fused_moe_max_dff_from` that the cap is family-agnostic (governs every SwitchGLU MoE model per the module `Used by:` list) and that the backend is resolved at runtime via `metal_is_available()` rather than a `cfg!(feature = "cuda")` switch, matching the fused kernel's own runtime dispatch (`run_fused_moe_two_kernel` picks the `cuda_kernel` port when `metal::is_available()` is false). - `src/models/switch_layers.rs`: added `fused_moe_dff_above_cap_declines_and_at_cap_dispatches`, a pure test pinning that `forward_fused_kernel` declines exactly when `dff > max_dff` at the cap boundary on both backends and under an explicit override. - `docs/benchmark_results/fused-moe-decode-kernel-design.md`: added the 2026-07-09 addendum (CUDA default provenance, why keeping 8192 is safe under pin drift, and the pending GB10 re-validation on MLX 0.32.1 with the exact `mlxcel-bench-decode` harness to run it), and cross-referenced it from the `MLXCEL_FUSED_MOE_MAX_DFF` table row. ## Design notes - The value 8192 is unchanged. Under pin drift the most conservative floor of the measured crossover (break-even at 8192, regression only at 14336) is the correct choice, so no re-tuning is done without hardware re-validation. - Metal behavior is byte-identical by construction (still 4096). The diff is runtime-inert (doc comments, a doc file, and a test only); no compiled non-test code path changed. ## Test plan Ran here (Apple M1 Ultra, Metal): - [x] `cargo fmt --all -- --check` - [x] `cargo test --release --features metal,accelerate -p mlxcel --lib switch_layers` (8/8 pass, including the new `fused_moe_dff_above_cap_declines_and_at_cap_dispatches`; the successful release build subsumes `cargo check --lib --tests`) - [x] `cargo clippy --features metal,accelerate -p mlxcel --lib --tests -- -D warnings` (clean) - [x] Metal MoE decode smoke, dispatch unchanged: `mlxcel generate -m models/qwen1.5-moe-a2.7b-4bit -p "Hello" -n 32 --temp 0` produced coherent output at 96.58 tok/s on Apple GPU (Metal), exercising the fused MoE decode path (Dff below the 4096 Metal cap, so it dispatches). Documented-data (not re-run in this PR): the CUDA default 8192 rests on the GB10 sweep in the 2026-07-03 addendum under MLX pin `e9463bb` (#626), landed in #643. Pending (hardware-blocked): GB10 re-validation of the crossover on the current MLX 0.32.1 pin. No CUDA/GB10 hardware was reachable for this change; the how-to-run harness is recorded in the 2026-07-09 addendum. Closes #330

## Summary Re-validates the CUDA fused-MoE Dff cap sweep behind `FUSED_MOE_MAX_DFF_CUDA` on GB10 under the current MLX 0.32.1 pin, closing the "pending re-validation" note left by issue #330 / PR #711 in docs/benchmark_results/fused-moe-decode-kernel-design.md. ## Sweep table Hardware: GB10 (DGX Spark), sm_121, CUDA 13.0. MLX pin: `57c66cac7cb3e5b1eb350488a61f1506b40d39f8` (0.32.1, #703/#704). Harness: `mlxcel-bench-decode`, prompt "Hello, how are you today?", 100 decode tokens after a 20-token warmup, median of 3 runs per side, `MLXCEL_FUSED_MOE_MAX_DFF=1` (gather_qmm fallback) vs `=20000` (fused). | Model | Dff | fallback tok/s | fused tok/s | delta | |-------|----:|----------------:|------------:|------:| | lfm2-8b-a1b (control) | 1792 | 140.68 | 160.63 | +14.2% | | phi-3.5-moe | 6400 | 53.85 | 55.80 | +3.6% | | llama-4-scout-17b | 8192 | 21.64 | 21.32 | -1.5% | | mixtral-8x7b | 14336 | 28.26 | 27.78 | -1.7% | The ratio=1.0 crossover interpolates to ~7672 between the phi-3.5-moe and llama-4-scout points, close to the prior ~8000 measured on MLX pin e9463bb. The ~4% shift sits within run-to-run noise at this Dff point: the 2026-07-03 sweep already flagged Dff 8192 as noisy (one of its three fused runs matched the fallback), and this 2026-07-10 sweep's own llama-4-scout fallback spans 21.33-21.85 tok/s run to run, a 2.4% spread comparable to the shift. ## What changed - `FUSED_MOE_MAX_DFF_CUDA` in `src/models/switch_layers.rs` is confirmed at 8192, not re-tuned; this PR does not touch that file. - Adds a 2026-07-10 addendum to docs/benchmark_results/fused-moe-decode-kernel-design.md with the sweep table, per-model findings, and the confirm-not-retune conclusion. - Amends the 2026-07-09 addendum's "pending" note to point at the new addendum, keeping the original text as historical provenance. - Updates the `MLXCEL_FUSED_MOE_MAX_DFF` env-var table row to state the confirmed result instead of "pending". - Adds the raw per-run sweep data as `benchmarks/cuda_gb10_issue712_fused_moe_dff_2026-07-10.csv`. ## Test plan - [x] `cargo check --features cuda --lib --tests` (pinned 1.93.1 toolchain) - clean, `Finished dev profile` with no errors or warnings. - [x] `cargo test --features cuda --lib switch_layers` - 8 passed, 0 failed (unchanged since the constant was not modified). - [x] Real GB10 hardware sweep: 24 `mlxcel-bench-decode` runs (4 models x 2 arms x 3 repeats), one process per run, rc=0 on every run. Closes #712

inureyes added type:chore Maintenance tasks (build, CI, etc.) priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers status:review Under review labels Jul 9, 2026

inureyes added status:done Completed and removed status:review Under review labels Jul 9, 2026

inureyes merged commit 6c53673 into main Jul 9, 2026
5 checks passed

inureyes deleted the update/issue-703-mlx-pin branch July 9, 2026 02:31

This was referenced Jul 9, 2026

perf(metal): close native NVFP4 prefill gap #705

Closed

perf(moe): record CUDA fused-MoE Dff cap provenance and decline test #711

Merged

perf(moe): re-validate the CUDA fused-MoE Dff cap sweep on MLX 0.32.1 (GB10) #712

Closed

inureyes mentioned this pull request Jul 9, 2026

docs(moe): re-validate CUDA fused-MoE Dff cap on MLX 0.32.1 #721

Merged

3 tasks

This was referenced Jul 19, 2026

fix(models): DeepSeek-V2-Lite emits nondeterministic garbage on GB10 CUDA (broken since the 2026-05-19 benchmark) #817

Closed

fix(models): dedicated root-cause debugging of the DeepSeek-V2-Lite greedy nondeterminism (follow-up to #817) #824

Closed

inureyes self-assigned this Aug 31, 2026

Navigation Menu

Uh oh!

chore: bump MLX pin to 0.32.1 (#703) - #704

chore: bump MLX pin to 0.32.1 (#703)#704
inureyes merged 2 commits into
mainfrom
update/issue-703-mlx-pin

inureyes commented Jul 9, 2026 •

edited

Loading

Uh oh!

inureyes commented Jul 9, 2026

Uh oh!

inureyes commented Jul 9, 2026

Uh oh!

inureyes commented Jul 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

inureyes commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

inureyes commented Jul 9, 2026

Implementation Review Summary

Intent

Findings Addressed

Remaining Items

Verification

Uh oh!

inureyes commented Jul 9, 2026

Uh oh!

inureyes commented Jul 9, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

inureyes commented Jul 9, 2026 •

edited

Loading