Repository navigation
bench: add M3 Ultra baselines for the three large MoE checkpoints - #2476
Merged
Merged
Conversation
bench_mlxlm.py hardcoded its memory budget at 85% of 128 GB, so on the 512 GB M3 Ultra it skipped deepseek-v4-flash-4bit and both inkling-small checkpoints as SKIP:oom_estimate even though bench_decode.sh had just measured them. The budget now follows the host's own memory (hw.memsize on macOS, MemTotal on Linux), mirroring bench_decode.sh, and falls back to 128 GB when memory cannot be read. PYLM_BENCH_MAX_GB and the ROCm carve-out override still apply. deepseek_v4 and inkling exist only in mlx-vlm, so the plain text path recorded FAIL:warmup for them. The new --text-via-vlm flag runs the same text protocol (synthesized 512-token prompt, no chat template, no image) but loads through mlx-vlm, records mlx-vlm as baseline_version, and writes to a separate pylm_*_textvlm_* file. The ten CSVs are the 2026-09-08 M3 Ultra measurements: mlxcel text and VLM rows plus the mlx-vlm text and VLM baselines for the same three checkpoints. Validated with py_compile, --help, and a module load that resolves the budget to 435 GB on this host.
inureyes
added
type:enhancement
The memory-budget fix added `_detect_memory_bytes()`, which shells out to `sysctl -n hw.memsize` on Darwin and parses `/proc/meminfo` on Linux. The file already had `_host_memory_bytes()`, a single portable `os.sysconf` call used by the ROCm hardware tagging, and both return the same value on the 512 GB M3 Ultra. The budget now calls `_host_memory_bytes()` and the duplicate is gone. The helper moves above the module-level budget computation, since the budget runs at import time and the helper used to be defined later in the file. Behavior is unchanged: 435.2 GiB on the 512 GB host, `PYLM_BENCH_MAX_GB` still overrides (65 GiB checked), and the ROCm path still resolves the helper.
The 2026-09-08 rows measured the pre-#2310 decode path, recorded two VLM rows on a dirty tree, and labeled the baseline with a stale 0.6.16 from a detached mlx-vlm pin. Swap all ten for a fresh 2026-10-20 four-way run at 058d7f7 on current main (MLX 243bdf09): mlxcel text and VLM rows plus mlx-vlm --text-via-vlm and VLM baselines for the same three checkpoints, baselines now recording mlx-vlm 0.7.9 from upstream main. The decode gap to the mlx-vlm baseline persists on M3 Ultra after #2310 wired the fused MoE kernel: mlxcel decodes at 23.7% (deepseek-v4-flash-4bit), 15.9% (inkling-small-mlx-4bit) and 18.5% (inkling-small-nvfp4) of the baseline on the text protocol, and 15.8% / 18.5% on the VLM protocol, each within ~3% of the September numbers on both sides. A same-build control with MLXCEL_FUSED_MOE=0 decodes at 7.80 tok/s against 7.78 with the fused kernel enabled (metal_m3ultra_2026-10-20_nofused_single_inkling-small-mlx-4bit.csv), so the kernel is not the lever on this host: the gap must come from elsewhere in the batch-1 expert path. The nvfp4 control arm stays on the gather path and moved by the same amount as the fused-path model, so the rerun is internally consistent. Refs #1858, #1859.
This was referenced Oct 20, 2026
inureyes
added a commit
that referenced
this pull request
Oct 20, 2026
…arriers (#2477) ## Summary Two CUDA overlay fixes for silently wrong values on GB10, each with tests that fail on the unfixed overlay, and both model-level workarounds removed: the OpenAI Privacy Filter sorts its dense expert planes again and the DeBERTa-v2 encoder (GLiNER2.5) no longer ends a CUDA graph after every layer. ## A1: sorted `gather_mm` read the stride of a size-1 axis as its leading dimension `gather_mm_rhs` (`matmul.cpp`, the `M == 1 && right_sorted_` branch every sorted MoE expert gather takes) passed `a.strides(-2)` to the grouped GEMM as the stride between consecutive gathered rows. There `a` is `[rows, 1, K]`, so that is the stride of a size-1 axis, which MLX ignores when it flags an array row-contiguous (`check_contiguity`) and leaves at whatever the producer wrote. The branch now makes `a` row-contiguous and uses `lda = K`, the batch stride, as `GatherQMM`'s grouped path already did. The producer in the Privacy Filter is the per-expert bias add in `ExpertLinear::forward`, `out + expand_dims(take(bias, ids), -2)`: `expand_dims` inserts a stride of 1, and a binary op inherits its second operand's layout when it donates that operand, which happens only when the first operand's cached buffer is more than 16 KiB over size (the cache hands out buffers up to two pages larger). So gate and up carried strides `[N, 1, 1]` or not depending on what the allocator had handed out, the compiled SwiGLU and the bf16 cast inherited them, and the down projection's sorted gather read `lda = 1`: the first row of every expert group right, every later row shifted by one element. That is PR #2460's drift (`2024-05-17` sliding from `private_date` to `private_phone` after the 8-bit model had run). A deferred per-plane capture showed exactly that signature (gate and up off by one bf16 ulp, `down` rows 1.. of each group off by up to 7.7e2). Also recorded: the overlay promotes `bf16 + f32` to bf16 (`overlays/cuda/mlx/dtype.cpp`), so "f32 activations" never reach this kernel as f32; the stride bug is dtype-independent. Exposure: every `SwitchLinear::forward(.., sorted = true)` caller with non-quantized planes reaches the branch (deepseek, exaone_moe, glm4_moe, glm4_moe_lite, gpt_oss, hunyuan_moe, ernie4_5_moe, llama4, qwen3_next, gemma4, phixtral, the SwitchGLU families, openai_privacy_filter). Their `sorted_x` is a `take` output with natural strides; the producers that reach the branch with a unit stride today are the biased `ExpertLinear` planes of GPT-OSS and the Privacy Filter, so a dense (bf16) GPT-OSS checkpoint was exposed the same way (quantized ones take `gather_qmm`). ## A2: captured child graphs race with concurrent kernel nodes The first `encode` of a fresh process returned wrong hidden states in about 40% of processes (7 of 20 on main with the split forced off), from a random layer each time, every replay correct. Allocation poisoning (0xFF on every allocation) changed nothing and produced no NaN. Bisection, 20 fresh processes each: graphs off 0, one op per graph 0, the whole encode in one graph 0, graph exec cache of 1 (fresh instantiation per graph) 10 of 10, program-order chaining of every node 0, chaining only edges that touch a stream-captured child graph (the cuBLASLt matmuls) 0, chaining only kernel-kernel edges 18, zeroing the whole cuBLASLt workspace inside the capture 0, workspace ceiling 0 (no split-K) 10. The child graphs hold `cutlass_80_simt_sgemm_*` with `cublasLt::splitKreduce_kernel` and a 240-byte workspace memset; their declared buffers and the dumped edges are correct. The resource that makes a cuBLASLt child graph wrong when an unrelated kernel node of the same graph runs alongside it is not identified. Fix (`device.cpp` overlay, new): a captured child graph is a barrier inside a graph. It joins every open branch added before it and every node added after it depends on it, on top of upstream's buffer-derived edges; kernel nodes between two child graphs stay concurrent. 0 of 30 fresh processes with the barrier, 5 of 10 with `MLXCEL_CUDA_GRAPH_CAPTURE_BARRIER=0` on the same binary. `MLXCEL_CUDA_GRAPHS_DOT_VERBOSE=1` adds node parameters to `MLX_SAVE_CUDA_GRAPHS_DOT_FILE` dumps. Both are in `docs/environment-variables.md`. Review pass (`5e29de95`): the two switches parse with the docs grammar (trimmed, case-insensitive, `0`/`false`/`off`/`no` and `1`/`true`/`on`/`yes`) via a new `mlxcel_env_flag` helper, since MLX's `atoi` env reader reads `on` as 0 and would have disabled the barrier by default; a `ConcurrentContext` batch whose nodes include a captured child graph makes its join node the barrier and clears the open-branch state afterwards, so cuBLASLt matmuls grouped under a concurrent section gate later nodes the same way single nodes do; and the per-encoder barrier map is leaked on purpose (`~CommandEncoder` runs `commit` and `clear_graph_state` during thread exit, after a `thread_local` with a destructor may be gone). Perf A/B (same binary, barrier on vs off, 3 rounds interleaved, `mlxcel generate --profile` on a 1767-token prompt): prefill qwen2.5-0.5b-bf16 on 3274 / 3415 / 3413 tok/s vs off 3421 / 3277 / 3260; llama-3.2-1b-4bit on 3896 / 4094 / 3847 vs off 3879 / 3929 / 3947, both within the run-to-run spread. Decode of 64 tokens after the same prompt, qwen2.5-0.5b-bf16 on 151.6 / 156.9 tok/s vs off 155.1 / 154.2 / 157.1 (the first "on" run, 97.5, was the cold process); llama-3.2-1b-4bit on 165.1 / 167.0 vs off 164.9 / 163.4 / 166.8. Decode runs gemv kernels, so no child graph is captured and the barrier is inert. The f32 encoder (`mlxcel extract`, gliner2.5-small-v1, 1400-character text, wall including load): on 1.25 / 1.26 / 1.25 s vs off 1.19 / 1.19 / 1.25 s, about +4%, the prefill GEMMs no longer overlap their neighbours. `mlxcel tag` on openai-privacy-filter-bf16: on 2.27 / 2.34 s vs off 2.31 / 2.19 / 2.35 s. MoE decode A/B (bench-decode, same binary, barrier on vs off, 4 interleaved rounds per arm): qwen3-moe-4bit on 76.46 vs off 75.96 tok/s, phi-3.5-moe-4bit on 43.21 vs off 42.61, paired deltas alternate sign and sit inside the ~1.5% host noise floor, so the barrier costs nothing measurable on MoE decode either. ## Tests and code - `src/models/switch_layers_sorted_gather_layout_tests.rs` and `_history_tests.rs` (new, helpers in `_test_support.rs`): `sorted_gather_mm_with_a_unit_stride_on_the_singleton_axis` and `sorted_gather_mm_after_a_biased_expert_add` (the Privacy Filter's op chain, donation forced) fail on the unfixed overlay (max |err| 1.19 and 5e-1 against an f64 reference, versus one bf16 ulp) and pass on the fix; `sorted_gather_mm_with_a_transposed_layout_plane` covers the row-major B instantiation; three history tests run the sorted gather before and after an 8-bit `gather_qmm` MoE, allocator churn and other sorted shapes and compare byte for byte (they pass before and after: allocation history alone never triggered it). - `src/models/deberta_v2_tests.rs::real_encoder_first_run_equals_replay` (new): the whole trunk on the 306-token fixture, first encode equals replay bit for bit, run in N fresh processes. - `openai_privacy_filter_layers.rs` drops `sort_routed` (every plane sorts at 64 or more rows, upstream's threshold); `moe_block_runs_both_gather_paths_identically` now compares the sorted and unsorted paths for real. `deberta_v2.rs` drops `graph_break_per_layer`. `mlxcel_core::array_strides` is new (post-eval strides only; before eval they are the shape's defaults). ## Verification (GB10 sm_121, `--profile test-fast --features cuda`, `MLX_ENABLE_TF32=0`, test binary under gpu-lock with `--test-threads=1`) - A1: `sorted_gather` 6 passed; on a control build with the unfixed `matmul.cpp` the unit-stride and biased-add tests fail and the other four pass. `privacy_filter_8bit` then `privacy_filter_spans` in one process (PR #2460's failing order): 3 of 3 fail on the control build, 3 of 3 pass on the fix. `privacy_filter_checkpoint_tests` 5 passed on the sorted dense path (the acceptance sentence in the bf16, 8-bit, bf16 order; the 4536-token chunk-invariance check; the log-prob and span fixtures). - A2: `real_encoder_first_run_equals_replay` 0 failures in 20 fresh processes; the same binary with the barrier disabled fails 5 of 10, and 13 of 20 with `CUDA_MODULE_LOADING=EAGER` and the barrier disabled, so lazy module loading is not the trigger. Revalidated after the review fixes: 20 fresh processes with the barrier, 0 failures. `deberta_v2` 8, `gliner2_5` 22, `token_classification` 86, `openai_privacy_filter` 10 passed. - `cargo fmt --all -- --check` and `cargo clippy --profile test-fast --features cuda -p mlxcel --lib --tests --no-deps -- -D warnings` clean. Full lib suite with `--test-threads=1`: 9373 passed, 0 failed, 166 ignored. - Full workspace gate on `5e29de95` (161 test binaries run in place under gpu-lock with `--test-threads=1`): build 0, clippy 0, doctests 0, contract set 0; 160/161 bins pass, `mllama_parity` failed once with the known 1-ulp `ragged_multi_image_states_concatenate_real_rows` flake and passes on rerun (8/8); issue #2467. After merging main (PRs #2469 to #2476, one overlay README conflict resolved with both additions), the same gate on the merged head `3ae5e2c1` gives the same result: build 0, clippy 0, doctests 0, contract set 0, 160/161 with the same mllama flake passing on rerun. - Not covered here: Metal (untouched) and other CUDA parts. ## Upstream report drafts (not filed) A1, `[CUDA] gather_mm with sorted indices reads the stride of a size-1 axis as lda`: `GatherMM::eval_gpu` routes `M == 1 && right_sorted_` to `gather_mm_rhs`, which passes `ensure_batch_contiguous(a)`'s `strides(-2)` as `lda` to `cutlass_grouped_gemm_unaligned`; `a` is `[..., 1, K]` there, so that is a size-1 axis stride that `check_contiguity` ignores and `expand_dims` or a donating binary op leaves at 1, and every gathered row after the first of a group then reads shifted memory. Repro: `a = mx.swapaxes(mx.random.normal((96, 640, 1)), 1, 2)` (strides `(640, 1, 1)`), `w` bf16 `(128, 640, 640)`, sorted expert ids, `mx.gather_mm(a.astype(mx.bfloat16), mx.swapaxes(w, -1, -2), rhs_indices=idx, sorted_indices=True)` versus `sorted_indices=False`: rows that are first in their group agree, every other row differs by O(1). Fix: make `a` row-contiguous and use `lda = a.shape(-1)`, as `GatherQMM` already does. A2, `[CUDA] stream-captured child graph nodes (cuBLASLt matmuls) race with concurrent kernel nodes of the same graph`: GB10 sm_121, driver 580.178.04, CUDA 13, MLX 243bdf09c, `MLX_ENABLE_TF32=0`, f32 weights (fastino/gliner2.5-small-v1, 12 DeBERTa-v2 layers, hidden 384, 306 tokens). About 40% of fresh processes return a first forward that differs from every later one by 0.1 to 0.6, from a random layer; the bisection is the list in the A2 section above. Each cuBLASLt call has its own workspace, registered inputs and outputs, and correct edges in the dumped graphs. Workaround: captured child graphs as barriers inside a graph. Repro: lablup/mlxcel `src/models/deberta_v2_tests.rs::real_encoder_first_run_equals_replay` in 20 fresh processes with `MLXCEL_CUDA_GRAPH_CAPTURE_BARRIER=0`. Technical report: `TECHNICAL_REPORTS/pr-2477-cuda-gather-mm-graph-barrier.en.md` (Korean `.ko.md`). Closes #2462
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
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
scripts/bench_mlxlm.py: the memory budget now follows the host's memory instead of a hardcoded 128 GB, matchingbench_decode.sh. On the 512 GB M3 Ultra the old cap skipped all three large checkpoints thatbench_decode.shhad measured. The budget reads the host size through_host_memory_bytes, reused as the single source (058d7f7). Closes fix(bench): bench_mlxlm.py hardcodes a 128 GB memory budget #1858.scripts/bench_mlxlm.py: new--text-via-vlmflag runs the text protocol through mlx-vlm for architectures mlx-lm does not implement (deepseek_v4,inkling). Results go topylm_*_textvlm_*files withbaseline_versionset to mlx-vlm.benchmarks/: ten CSVs from a 2026-10-20 M3 Ultra four-way run at 058d7f7 on current main (mlxcel text/VLM rows plus mlx-vlm--text-via-vlmand VLM baselines for deepseek-v4-flash-4bit, inkling-small-mlx-4bit, inkling-small-nvfp4), plus one same-build control CSV. This replaces the 2026-09-08 set, which measured the pre-perf(moe): attribute M = 1 expert gather gap and wire Inkling/V4 MoE fast paths (#1859) #2310 decode path, recorded two rows on a dirty tree, and labeled the baseline with a stale 0.6.16 from a detached mlx-vlm pin.Decode gap after #2310 on M3 Ultra
#2310 wired the fused MoE kernel (inkling) and removed the LimitedSwitchGlu redundant graph nodes (deepseek_v4); this rerun is the first decode measurement on any host after that fix. The gap to the mlx-vlm baseline persists: on the text protocol mlxcel decodes at 23.7% (deepseek-v4-flash-4bit), 15.9% (inkling-small-mlx-4bit) and 18.5% (inkling-small-nvfp4) of the baseline, and on the VLM protocol at 15.8% / 18.5%, each within ~3% of the September numbers on both sides.
A same-build control with
MLXCEL_FUSED_MOE=0decodes at 7.80 tok/s against 7.78 with the fused kernel enabled (metal_m3ultra_2026-10-20_nofused_single_inkling-small-mlx-4bit.csv), so the kernel is not the lever on this host: the gap must come from elsewhere in the batch-1 expert path. inkling-small-nvfp4 stays on the gather path as the control arm and moved by the same amount as the fused-path model, so the rerun is internally consistent.Validation
python3 -m py_compile scripts/bench_mlxlm.pyPYLM_BENCH_MAX_GBstill overridesmlxcel_commit=058d7f78on a clean tree, MLX243bdf09