[Metal] global scale for qmm by nastya236 · Pull Request #4458 · ml-explore/mlx

nastya236 · 2026-09-03T23:29:08Z

This pull request adds global scales for MoE weights that are quantized to nvfp4 only for Metal.

Two important statements:

gather_qmm is used primarily on Metal for quantized MoEs.
nvfp4 format consists of 2 scales: fp8_block_scale for 16 elements and fp32_global_scale per tensor scale.
Then to dequantize the weight W fo 16 elements block w:
w = e2m1_code × fp8_block_scale × fp32_global_scale

Naive implementation would be to multiply the result of the multiplication:

x = x @ dequant(W) * fp32_global_scale

which is incorrect, because fp32_global_scale should be used in dequantization kernel to get a correct result

Note: it looks like a big diff but it is just adding new parameter to all related operations + changing dequantization part in gather_qmm.
Example:

import mlx.core as mx

K = 4096

q_fp4 = mx.full((K, 1), 6.0, dtype=mx.float32)
s_fp8 = mx.array(0.5, dtype=mx.float32)
s_dec = mx.array(0.1, dtype=mx.float32)

x = mx.ones((1, K), dtype=mx.bfloat16)

# dequantize without global scale, then scale output
W_without_global = (q_fp4 * s_fp8).astype(mx.bfloat16)
y_scale_after = (x @ W_without_global) * s_dec

# apply complete scale before converting weight to BF16
W_with_global = (q_fp4 * s_fp8 * s_dec).astype(mx.bfloat16)
y_scale_inside = x @ W_with_global

print("W without global:", W_without_global[0, 0].item())
print("W with global:", W_with_global[0, 0].item())
print("Scale after:", y_scale_after.item())
print("Scale inside:", y_scale_inside.item())
print("Difference:", (y_scale_after - y_scale_inside).item())

W without global: 3.0
W with global: 0.30078125
Scale after: 1228.800048828125
Scale inside: 1232.0
Difference: -3.199951171875

TODO: add the same for CUDA for gather_qmv.

zcbenz

Nice work! just some nitpickings.

zcbenz · 2026-09-07T09:42:09Z

+                                bn,
+                                wm,
+                                wn,
+                                true)


I think you can simplify the big if clause with:

kernel = get_qmm_nax_kernel_wrapped( ... global_scale.has_value());

You right! I just added global scale always false to affine quantization.

zcbenz · 2026-09-07T09:46:21Z

+            ).reshape((E, 1, 1))
+
+            for dtype in [mx.float32, mx.float16, mx.bfloat16]:
+                for transpose in [True, False]:


You can reduce the indentations by using product, for example:

mlx/python/tests/test_quantized.py

Lines 230 to 235 in ce916db

tests = product(

[256, 512, 67], # M

[64, 256, 512], # N

["nvfp4", "mxfp8"], # mode

)

for M, N, mode in tests:

…al-scale

) ## 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 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` (ml-explore/mlx#4458), so all 13 calls pass `std::nullopt`. - CUDA now needs 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 (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

nastya236 added 2 commits September 4, 2026 01:10

added global scale for qmm

727c58e

fix metal::is_available

f7c3404

nastya236 force-pushed the gather-qmm-nvfp4-global-scale branch from 8a420c9 to f7c3404 Compare September 3, 2026 23:45

nastya236 added 2 commits September 4, 2026 01:46

Merge branch 'main' into gather-qmm-nvfp4-global-scale

d6f1a8e

fix shapes and errors

285ab89

nastya236 mentioned this pull request Sep 4, 2026

Add global scale for nvfp4 quantized MoEs ml-explore/mlx-lm#1844

Open

zcbenz approved these changes Sep 7, 2026

View reviewed changes

nastya236 added 2 commits September 8, 2026 14:54

Merge remote-tracking branch 'origin/main' into gather-qmm-nvfp4-glob…

a54b48a

…al-scale

address comments

af85e41

nastya236 merged commit cbb4059 into main Sep 8, 2026
29 checks passed

nastya236 deleted the gather-qmm-nvfp4-global-scale branch September 8, 2026 14:40

inureyes mentioned this pull request Sep 10, 2026

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

Merged

aaishwarymishra pushed a commit to aaishwarymishra/mlx that referenced this pull request Sep 11, 2026

[Metal] global scale for qmm (ml-explore#4458)

ee7eaee

nastya236 mentioned this pull request Sep 16, 2026

[CUDA] Add global scale support to gather_qmm #4507

Merged

davidtai mentioned this pull request Sep 28, 2026

[Metal] gather_mm and gather_qmm row-tile scheduling (backport of ml-explore/mlx #4567, #4572) Layr-Labs/mlx#27

Merged

Pushkinist mentioned this pull request Sep 30, 2026

MLX 0.32: cross-thread eval throws; move rMLX to a self-built mlx 0.32.3 / mlx-c 0.7.0 with NAX Pushkinist/rMLX#638

Open

2 tasks done

Navigation Menu

Repository navigation

[Metal] global scale for qmm - #4458

[Metal] global scale for qmm#4458
nastya236 merged 6 commits into
mainfrom
gather-qmm-nvfp4-global-scale

nastya236 commented Sep 3, 2026 •

edited

Loading

Uh oh!

zcbenz left a comment

Uh oh!

zcbenz Sep 7, 2026

Uh oh!

nastya236 Sep 8, 2026

Uh oh!

zcbenz Sep 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	tests = product(
	[256, 512, 67], # M
	[64, 256, 512], # N
	["nvfp4", "mxfp8"], # mode
	)
	for M, N, mode in tests:

Conversation

nastya236 commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zcbenz left a comment

Choose a reason for hiding this comment

Uh oh!

zcbenz Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

nastya236 Sep 8, 2026

Choose a reason for hiding this comment

Uh oh!

zcbenz Sep 7, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

nastya236 commented Sep 3, 2026 •

edited

Loading