Skip to content

b10994

Pre-release
Pre-release

Choose a tag to compare

View all tags
@github-actions github-actions released this 16 Sep 07:23
b10994
0a8b29a
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

metal: fix NaN in mul_mm_id when activations exceed f16 range (#26223)

  • test-backend-ops: reproduce MUL_MAT_ID NaN for activations beyond f16

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).

  • metal: fix NaN in mul_mm_id when activations exceed f16 range

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 #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). #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).

  • metal : make requested edits
  • remove verbose comments
  • explain rationale as requested

Generative AI disclosure: Claude made the edits as requested.

  • metal : stack mul_mm_id map0 with amax_part

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

  • cont : fix var scope

  • cont : comment out tests temporarily

Comment out tess to not break CI temporarily

Assisted-by: Claude Fable 5.1

Co-authored-by: Claude Fable 5 noreply@anthropic.com
Co-authored-by: Georgi Gerganov ggerganov@gmail.com

Website:

Attestations:

macOS/iOS:

Linux:

Android:

Windows:

openEuler:

  • DISABLED
  • openEuler x86 (310p)
  • openEuler x86 (910b, ACL Graph)
  • openEuler aarch64 (310p)
  • openEuler aarch64 (910b, ACL Graph)

UI:

Contributors

ggerganov
Assets 35
Loading