metal: fix NaN in mul_mm_id when activations exceed f16 range by mdegans · Pull Request #26223 · ggml-org/llama.cpp

mdegans · 2026-07-28T10:32:37Z

Fixes an all-NaN output from mul_mm_id on Metal when a model's activations exceed f16 range. Likely fixes #25722; #20668 may be the same defect attributed to a bad GGUF.

The bug

kernel_mul_mm_id narrows src1 to half for the simdgroup MMA operands — S1 = half in every instantiation. f16 saturates at 65504, so activations above that become inf on load, and simdgroup_multiply_accumulate then propagates NaN across the entire 8x8 accumulator tile. The output is not degraded, it is entirely NaN.

The conversion sites are ggml-metal.metal in kernel_mul_mm_id:

*(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
*(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));

(plus the two mirrors in the GGML_METAL_HAS_TENSOR path)

Because the mul_mv_id path taken below ne21_mm_id_min (32) keeps the same values in f32, the same model produces correct logits for short prompts and NaN for long ones, with the switch at exactly 32 tokens. Accumulation was already f32 and was never the problem — only the operand narrowing.

Reproducer

The first commit adds one. This class of bug was previously untestable: init_mul_mat_id_tensors initializes uniform [-1, 1], so no existing case can drive an operand out of f16 range. test_mul_mat_id gains an amax parameter (default 1.0f, preserving the historical init exactly) that scales only the f32 activations, leaving quantized weights in their normal range.

Six cases, two shapes — n=16 is the control on the mul_mv_id path and must stay green; n=32/n=64 are above the switch:

MUL_MAT_ID(type_a=q8_0,type_b=f32,n_mats=8,n_used=2,b=0,m=512,n=32,k=256,amax=100000.000000):
  [MUL_MAT_ID] NaN at index 0 (MTL0=nan CPU=583442.375000) FAIL

It reproduces at minimal size (8 experts, 2 active, 512x256), so this is not specific to any model or geometry.

The fix

Rescale src1 by a power of two so it fits, and undo the scale on the f32 accumulator at the store. A two-stage reduction computes max(|src1|) and writes (1/scale, scale) into scratch chained off the destination buffer, in the same style as the existing tpe/ids id-mapping scratch.

This is exact, not approximate, for two reasons: the dot product is linear, so one tensor-wide factor commutes through the accumulation; and the factor is a power of two, so both multiplications are exact in binary floating point. When max(|src1|) already fits — every model that works today — the factor is exactly 1.0 and the output is bit-identical to before.

The reduction is two-stage (256 threadgroups into partials, then one threadgroup folding them) specifically so it stays bandwidth-bound; a single-threadgroup version was measured first and cost up to +451% on prefill. It is dispatched only on the mm path, so decode never pays for it.

Performance

Apple M2 Max, test-backend-ops perf -o MUL_MAT_ID -b MTL0, 99 cases, against the same build without this change:

n	path	median
1 / 4 / 8	`mul_mv_id` (decode)	-0.76% / -0.84% / -0.39% (noise)
32	`mul_mm_id` (prefill)	+1.73%
64		+1.30%
128		+1.80%
256		+3.98%
512		+3.74% (worst case +7.20%)
overall		+1.14%

Validation

the six new cases go from 4 FAIL / 2 OK to all OK, n=16 controls unchanged
test-backend-ops -b MTL0 full run: 0 failures, no regression
Mistral-Small-4-119B (arch mistral4, 128 experts / 4 active) generates correctly at the default n_ubatch of 512 in both UD-IQ3_S and UD-Q4_K_XL. Before this, every prefill of >=32 tokens returned an all-NaN vocabulary, and only n_ubatch <= 31 — forcing the mul_mv_id path — worked.

Not covered

kernel_mul_mm (dense) has the identical narrowing at the corresponding load sites and is expected to fail the same way. Left alone here to keep this reviewable; happy to extend if you'd prefer one change.
The scale could be per output column rather than per tensor, which would preserve more precision when a single token is the hot one. Not needed for the bug.

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES

This bug was found, diagnosed, reproduced and fixed by Claude Opus 5 (via Claude Code), working from a real-model failure. I reviewed and understand the code.

🤖 Generated with Claude Code

mdegans · 2026-07-28T12:17:46Z

AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

Noted for next time. I reviewed the post and was satisfied so I said "go for it".

ggerganov · 2026-07-28T17:30:30Z

+    // 16 bytes for the factor pair (inverse scale applied to src1 on
+    // load, scale applied to the f32 accumulator on store), then one
+    // float per stage-1 partial.


Shorten to:

Suggested change

// 16 bytes for the factor pair (inverse scale applied to src1 on

// load, scale applied to the f32 accumulator on store), then one

// float per stage-1 partial.

// 2 scaling factors (16 bytes) + N_MM_NPART_AMAX per-threadgroup scales for stage-1

Btw, why do we reserve 16 bytes for the 2 floats, instead of just 8 bytes?

ggerganov · 2026-07-28T17:31:16Z

+        // src1 rescale factors, computed before the matmul so the
+        // narrowing to the half MMA operands cannot overflow. See
+        // kernel_mul_mm_id_amax_f32.


Suggested change

// src1 rescale factors, computed before the matmul so the

// narrowing to the half MMA operands cannot overflow. See

// kernel_mul_mm_id_amax_f32.

// src1 rescale factors, computed before the matmul

// ref: https://github.com/ggml-org/llama.cpp/pull/26223

ggerganov · 2026-07-28T17:31:51Z

+// threadgroups into partials, stage 2 folds the partials and writes the
+// factor pair.
+
+#define GGML_METAL_AMAX_NPART 256


Rename this constant to N_MM_NPART_AMAX and move to ggml-metal-impl.h

ggerganov · 2026-07-28T17:32:55Z

+// Compute max(|x|) over src1 and derive a power-of-two scale that
+// brings the activations inside f16 range.
+//
+// kernel_mul_mm_id narrows src1 to `half` for the simdgroup MMA
+// operands. f16 saturates at 65504, so a model whose activations exceed
+// that yields inf, and simdgroup_multiply_accumulate then poisons the
+// whole accumulator tile with NaN. Scaling src1 down by a power of two
+// on load and scaling the f32 accumulator back up on store is exact —
+// powers of two are exact in binary floating point and the dot product
+// is linear, so a single tensor-wide factor introduces no error at all.
+//
+// When amax fits already (the overwhelmingly common case) the factor is
+// 1.0 and the result is bit-identical to not doing this.
+//
+// Two stages so the pass stays bandwidth-bound rather than serialized
+// on one threadgroup: stage 1 reduces rows across GGML_METAL_AMAX_NPART
+// threadgroups into partials, stage 2 folds the partials and writes the
+// factor pair.


Delete this comment - we have a reference to this PR at the kernel launch site.

ggerganov · 2026-07-28T17:33:19Z

+        // operand itself has to fit.
+        float scale = 1.0f;
+
+        if (isfinite(amax) && amax > 32768.0f) {


Is this isfinite check needed?

ggerganov · 2026-07-28T17:34:19Z

+    // Power-of-two rescale so activations outside f16 range survive the
+    // narrowing to the `half` MMA operands. Both factors are exactly 1.0
+    // unless src1 needed it, in which case they are exact powers of two.


Suggested change

// Power-of-two rescale so activations outside f16 range survive the

// narrowing to the `half` MMA operands. Both factors are exactly 1.0

// unless src1 needed it, in which case they are exact powers of two.

// power-of-two rescaling

ggerganov · 2026-07-28T17:35:08Z

+    // Activations outside f16 range. Backends that narrow src1 to a
+    // half-precision type for a matrix-multiply path (Metal's
+    // mul_mm_id feeds simdgroup_half8x8) saturate at 65504 and produce
+    // inf, then NaN — while the vector path on the same backend, and
+    // every CPU path, are correct. Real models hit this: Mistral
+    // Small 4 (arch mistral4, 128 experts / 4 active) has a layer whose
+    // ffn activations reach ~1e5, so on Metal every prefill of >= 32
+    // tokens returns an all-NaN vocabulary, and fewer than 32 tokens is
+    // correct (ne21_mm_id_min switches mul_mv_id -> mul_mm_id at 32).
+    // n = 32 and 64 sit above that switch; n = 16 below it is the
+    // control that must stay green.


Suggested change

// Activations outside f16 range. Backends that narrow src1 to a

// half-precision type for a matrix-multiply path (Metal's

// mul_mm_id feeds simdgroup_half8x8) saturate at 65504 and produce

// inf, then NaN — while the vector path on the same backend, and

// every CPU path, are correct. Real models hit this: Mistral

// Small 4 (arch mistral4, 128 experts / 4 active) has a layer whose

// ffn activations reach ~1e5, so on Metal every prefill of >= 32

// tokens returns an all-NaN vocabulary, and fewer than 32 tokens is

// correct (ne21_mm_id_min switches mul_mv_id -> mul_mm_id at 32).

// n = 32 and 64 sit above that switch; n = 16 below it is the

// control that must stay green.

// test src1 f16 overflow

ggerganov · 2026-07-28T17:35:42Z

+    // Magnitude of the src1 activations. Default 1.0f reproduces the
+    // historical uniform [-1, 1] init. Larger values exercise backends
+    // that narrow the activations to a lower-range type internally:
+    // Metal's mul_mm_id feeds simdgroup_half8x8, and f16 saturates at
+    // 65504, so real models whose activations exceed that produce inf
+    // and then NaN on that path while the mul_mv_id path is correct.
+    const float amax;


Suggested change

// Magnitude of the src1 activations. Default 1.0f reproduces the

// historical uniform [-1, 1] init. Larger values exercise backends

// that narrow the activations to a lower-range type internally:

// Metal's mul_mm_id feeds simdgroup_half8x8, and f16 saturates at

// 65504, so real models whose activations exceed that produce inf

// and then NaN on that path while the mul_mv_id path is correct.

const float amax;

const float amax; // magnitude of src1

ggerganov · 2026-07-28T17:36:25Z

+            // src1 (activations) only — the weights stay in their normal
+            // range so this isolates activation magnitude.


Delete the comment

ggerganov · 2026-07-28T17:36:53Z

+        // Leave a comfortable margin below the f16 max of 65504: the
+        // products feeding the accumulator stay in f32, so only the
+        // operand itself has to fit.


Suggested change

// Leave a comfortable margin below the f16 max of 65504: the

// products feeding the accumulator stay in f32, so only the

// operand itself has to fit.

// leave a comfortable margin below the f16 max of 65504

mdegans · 2026-07-29T10:00:27Z

@ggerganov

Would you like me to trim down the PR description as well? My justification for leaving it was/is that it explains the rationale of the agent who made the edits, however if you'd prefer one or two sentences I can make that happen.

ggerganov · 2026-07-29T12:26:33Z

+
+        ggml_metal_op_concurrency_reset(ctx);
+


In terms of concurrency, currently we run the kernels like this:

amax_part amax + map0 (these 2 run in parallel) main

I wonder if it would be more optimal to stack the map0 together with the amax_part like this:

amax_part + map0 (these 2 run in parallel) amax main

Ran ABBA BAAB tests on a hot M2 MacBook. B (amax_part + map0) is slightly faster (-1.8% mean, up to 3.5%) when the batch size is 512. At 32 it's just noise. If you'd like I can submit B as a separate PR. I'm assuming you don't want it rolled into this one.

It's ok to push the change in this PR. Btw, it's worth benchmarking up to -ub 2048 since many models benefit from larger than 512 microbatch size.

My mac's GPU is going to be tied up until Saturday. I will test up to 2048, after are reboot, then and add the commit here unless there is a performance regression.

So. As you suspected, win scales with batch size. The commit is added. The 1024 and 2048 cases are left out of the commit since everywhere else seems to only go to 512.

ggerganov · 2026-07-29T12:27:54Z

@ggerganov

Would you like me to trim down the PR description as well? My justification for leaving it was/is that it explains the rationale of the agent who made the edits, however if you'd prefer one or two sentences I can make that happen.

The description in the PR is OK. The goal is the comments in the code to be short to make it easier to read/edit. Longer descriptions should be referenced with links.

The Metal mul_mm_id path narrows src1 to `half` for the simdgroup MMA (`S1 = half` in every instantiation; ggml-metal.metal:10582 and :10595, mirrored at :10643/:10654 in the tensor-ops path). f16 saturates at 65504, so a model whose activations exceed that produces inf, and `simdgroup_multiply_accumulate` then turns the whole 8x8 accumulator tile into NaN. The mul_mv_id path used below `ne21_mm_id_min` (32) carries the same values in f32 and is correct, as is every CPU path. This was untestable before: `init_mul_mat_id_tensors` initializes uniform [-1, 1], so no existing case can drive an operand out of f16 range. `test_mul_mat_id` gains an `amax` parameter (default 1.0f, preserving the historical init exactly) that scales only the f32 activations, leaving the quantized weights in their normal range. Six cases: n=16 sits below the mul_mv_id -> mul_mm_id switch and is the control that must stay green; n=32 and n=64 are above it and fail on Metal today. Two shapes, because this is not model- or size-specific — q4_K at 128 experts / 4 active / 4096x2048 mirrors a real model, and q8_0 at 8 experts / 2 active / 512x256 shows the same failure at minimal size. Observed on Apple M2 Max, macOS, llama.cpp b10156: MUL_MAT_ID(type_a=q8_0,...,n=32,k=256,amax=100000.000000): [MUL_MAT_ID] NaN at index 0 (MTL0=nan CPU=583442.375000) FAIL The real model behind this is Mistral Small 4 (arch mistral4, 128 experts / 4 active), one of whose layers reaches ~1e5 activations: on Metal every prefill of >=32 tokens returns an entirely NaN vocabulary, while <32 tokens is correct. Note kernel_mul_mm (dense) has the identical conversion at :10273 and :10286 and is expected to fail the same way; it is not covered here. Found and written by Claude Opus 5 (via Claude Code).

kernel_mul_mm_id narrows src1 to `half` for the simdgroup MMA operands (`S1 = half` in every instantiation). f16 saturates at 65504, so a model whose activations exceed that produces inf on load, and simdgroup_multiply_accumulate then propagates NaN across the whole 8x8 accumulator tile. The result is an entirely NaN output — not a precision loss, a total loss. The mul_mv_id path taken below ne21_mm_id_min (32) keeps the same values in f32 and is correct, as is every CPU path, so the same model produces correct logits for short inputs and NaN for long ones. Fix: rescale src1 by a power of two so it fits, and undo the scale on the f32 accumulator at the store. A two-stage reduction computes max(|src1|) and writes the pair (1/scale, scale) into scratch chained off the destination buffer, in the same style as the existing tpe/ids id-mapping scratch. The matmul multiplies on load and on store. This is exact, not approximate, for two reasons: the dot product is linear, so one tensor-wide factor commutes through the accumulation; and the factor is a power of two, so both multiplications are exact in binary floating point. When max(|src1|) already fits — every model that works today — the factor is exactly 1.0 and the output is bit-identical to before. Accumulation was already f32 and is unchanged; only the operand narrowing was ever the problem. The reduction is two-stage (256 threadgroups into partials, then one threadgroup folding them) specifically so it stays bandwidth-bound. A single-threadgroup version was measured first and cost up to +451% median on prefill — the scan serialized against an otherwise idle GPU. It is also dispatched only on the mm path, so decode never pays for it. Measured on Apple M2 Max, `test-backend-ops perf -o MUL_MAT_ID -b MTL0`, 99 cases, versus the same build without this change: n=1/4/8 (mul_mv_id, decode) : -0.8% / -0.8% / -0.4% median (noise) n=32 (mul_mm_id, prefill) : +1.73% median n=64 : +1.30% median n=128 : +1.80% median n=256 : +3.98% median n=512 : +3.74% median, +7.20% worst overall : +1.14% median Correctness, same machine: - the six new test-backend-ops cases go from 4 FAIL / 2 OK to all OK, with the n=16 controls (mul_mv_id path) unchanged; - `test-backend-ops -b MTL0` full run: 0 failures, no regression; - Mistral-Small-4-119B (arch mistral4, 128 experts / 4 active) now generates correctly at the default n_ubatch of 512, in both UD-IQ3_S and UD-Q4_K_XL quantizations. Before this, every prefill of >= 32 tokens returned an all-NaN vocabulary and only n_ubatch <= 31 (forcing the mul_mv_id path) worked. Likely fixes ggml-org#25722 (mistral4 empty output on Metal above ~300 tokens, FA on and off, generation degenerating to a single control token — the signature of argmax over an all-NaN distribution). ggml-org#20668 may be the same defect attributed to a bad GGUF. Note kernel_mul_mm (dense) has the identical narrowing at the corresponding load sites and is expected to fail the same way; it is left alone here to keep this change reviewable. Also possible, and left for later: scaling per output column rather than per tensor, which would preserve more precision when a single token is the hot one. Found, diagnosed and fixed by Claude Opus 5 (via Claude Code).

- remove verbose comments - explain rationale as requested Generative AI disclosure: Claude made the edits as requested.

ggerganov · 2026-08-02T04:55:21Z

+
+            ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 32, 1, 1);
+        }
+
        // this barrier is always needed because the next kernel has to wait for the id maps to be computed


This comment should become:

// the next kernel has to wait for the amax data

It's done. Let me know when/if you want me to rebase on top of the latest master.

@ggerganov

Implement @ggerganov suggestion to stack amax_part + map0. Mean 2.6% faster (worst -0.7%, best -4.1%). Win grows with batch size. Benchmarked on a hot M2 Max after reboot. Generative AI disclosure: Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

kwakayama · 2026-08-22T07:32:03Z

Verified this fixes the failure on a second machine and a different chip.

Setup: Apple M4 Max / 128 GB, macOS, Metal. Built from 2100e5926 with -DGGML_METAL=ON.

Red/green with this PR's own test cases, toggling only ggml/ and rebuilding in between:

`ggml/`	result
reverted	`FAIL` — 4 NaN, 2/3 backends passed
applied	`OK` — 0 NaN, 3/3 backends passed

The n sweep behaves exactly as the description predicts:

MUL_MAT_ID(type_a=q4_K,...,n=16,...,amax=100000): OK      <- mul_mv_id path, control
MUL_MAT_ID(type_a=q8_0,...,n=32,...,amax=100000): FAIL
MUL_MAT_ID(type_a=q4_K,...,n=64,...,amax=100000): FAIL

n=16 stays green throughout, so the failure tracks ne21_mm_id_min rather than anything about the shapes.

Independent corroboration at the model level. Before finding this PR I hit #25722 with Mistral Small 4 119B (Q4_K_M) and bisected it from the other direction — via -ub on llama-server, holding the prompt fixed at 619 tokens:

`-ub`	output
31	`2+2 equals 4.`
32	(empty)

Same integer, reached without touching the kernel. That's consistent with your analysis that mul_mv_id below ne21_mm_id_min keeps the values in f32: capping n_ubatch at 31 means the mul_mm_id path is simply never reached, and arbitrarily long prompts then work.

One note for anyone else trying to reproduce the test: the amax scaling applies to every f32 tensor, so the quantized type_a in these cases is load-bearing. I first wrote equivalent cases with f32/f16 weights and got CPU=nan in the reference — both operands were being scaled. Using quantized weights keeps the scaling on the activations only, which is presumably why you chose them.

ggerganov · 2026-08-27T07:07:50Z

@mdegans Could you rebase on latest master.

@forforever73 I think this change is good. Could you have a second look when you get the chance? One improvement we can do in the future is to gate it with the new ggml_prec logic, but for now it should be fine to use it unconditionally.

Confirmed via a CPU-only diagnostic (-ngl 0) that the original empty-content/hallucination failure is a Metal-backend f16 overflow in mul_mm_id's MoE down-projection (large Laguna activations overflow f16's 65504 ceiling -> NaN), not a reasoning-mode or harness-probe issue. Upstream fix (ggml-org/llama.cpp#26223) is still unmerged. 9/9 clean results with GPU offload disabled entirely.

mdegans · 2026-08-30T01:20:26Z

@mdegans Could you rebase on latest master.

@forforever73 I think this change is good. Could you have a second look when you get the chance? One improvement we can do in the future is to gate it with the new ggml_prec logic, but for it should be fine to use it unconditionally.

I have been vacation but I can certainly take care of that in a few days when I am back home.

mdegans requested review from a team and ggerganov as code owners July 28, 2026 10:32

This comment was marked as resolved.

Sign in to view

github-actions Bot added testing Everything test related ggml changes relating to the ggml tensor library for machine learning Apple Metal https://en.wikipedia.org/wiki/Metal_(API) labels Jul 28, 2026

mdegans mentioned this pull request Jul 28, 2026

Prediction iterators are infallible: decode errors become a panic loop mdegans/drama_llama#92

Open

ggerganov self-assigned this Jul 28, 2026

ggerganov reviewed Jul 28, 2026

View reviewed changes

mdegans force-pushed the fix/metal-mul-mm-id-f16-overflow branch from a26851f to fe448b6 Compare July 29, 2026 10:18

ggerganov reviewed Jul 29, 2026

View reviewed changes

mdegans added 3 commits August 1, 2026 15:03

metal : make requested edits

6854b4b

- remove verbose comments - explain rationale as requested Generative AI disclosure: Claude made the edits as requested.

mdegans force-pushed the fix/metal-mul-mm-id-f16-overflow branch from fe448b6 to 976221d Compare August 1, 2026 13:06

ggerganov reviewed Aug 2, 2026

View reviewed changes

mdegans force-pushed the fix/metal-mul-mm-id-f16-overflow branch from 976221d to 2f7259a Compare August 4, 2026 11:55

ggerganov mentioned this pull request Aug 6, 2026

llama-graph: prevent f16 overflow in MoE down-projection on Metal #25442

Closed

ggerganov added 2 commits August 20, 2026 14:15

cont : fix var scope

bf96f5b

Merge branch 'master' into pr/26223

e6a3398

kwakayama mentioned this pull request Aug 22, 2026

Eval bug: mistral4 empty output on Metal for prompts over ~300 tokens (clean GGUF, with and without flash attention) #25722

Open

		// src1 (activations) only — the weights stay in their normal
		// range so this isolates activation magnitude.

Uh oh!

Conversation

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

Uh oh!

The bug

Reproducer

The fix

Performance

Validation

Not covered

Requirements

Uh oh!

This comment was marked as resolved.

mdegans commented Jul 28, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mdegans commented Jul 29, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mdegans Aug 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ggerganov commented Jul 29, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kwakayama commented Aug 22, 2026

Uh oh!

ggerganov commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mdegans commented Aug 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mdegans commented Jul 28, 2026 •

edited

Loading

mdegans Jul 29, 2026 •

edited

Loading

mdegans Aug 1, 2026 •

edited

Loading

ggerganov commented Aug 27, 2026 •

edited

Loading