llama.cpp

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

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>