Commit 0a8b29a60 for llama.cpp

commit 0a8b29a607604625b4351760a849a96140464abe
Author: Michael de Gans <michael.john.degans@gmail.com>
Date:   Wed Sep 16 08:37:40 2026 +0200

    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>

diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index bf3d07e78..b510cb957 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -1063,6 +1063,40 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
     return res;
 }

+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib) {
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_mul_mm_id_amax_part_f32");
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    res.smem = 32*sizeof(float);
+
+    return res;
+}
+
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib) {
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_mul_mm_id_amax_f32");
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    res.smem = 32*sizeof(float);
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(ggml_metal_library_t lib, int ne02, int ne20) {
     char base[256];
     char name[256];
diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h
index ced33aadf..f6243ffbd 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.h
+++ b/ggml/src/ggml-metal/ggml-metal-device.h
@@ -138,6 +138,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ex
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0    (ggml_metal_library_t lib, int ne02, int ne20);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax(ggml_metal_library_t lib);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_amax_part(ggml_metal_library_t lib);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax            (ggml_metal_library_t lib, const struct ggml_tensor * op);
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 7ad21341e..7a2c65aaa 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -14,6 +14,8 @@
 #define N_MM_SIMD_GROUP_X 2
 #define N_MM_SIMD_GROUP_Y 2

+#define N_MM_NPART_AMAX 256
+
 // kernel parameters for mat-vec threadgroups
 //
 // N_R0: number of src0 rows to process per simdgroup
@@ -555,6 +557,14 @@ typedef struct {
     uint64_t nb21;
 } ggml_metal_kargs_mul_mm_id_map0;

+typedef struct {
+    int32_t  ne00;
+    int32_t  ne01;
+    int32_t  ne02;
+    uint64_t nb01;
+    uint64_t nb02;
+} ggml_metal_kargs_mul_mm_id_amax;
+
 typedef struct {
     int32_t  ne00;
     int32_t  ne02;
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index da0040a0c..cc1bebfaa 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -2631,6 +2631,15 @@ size_t ggml_metal_op_mul_mat_id_extra_ids(const ggml_tensor * op) {
     return ggml_type_size(GGML_TYPE_I32)*ne02*ne21;
 }

+size_t ggml_metal_op_mul_mat_id_extra_amax(const ggml_tensor * op) {
+    assert(op->op == GGML_OP_MUL_MAT_ID);
+
+    GGML_UNUSED(op);
+
+    // 2 scaling factors (8 bytes) + N_MM_NPART_AMAX per-threadgroup scales for stage-1
+    return 8 + N_MM_NPART_AMAX*sizeof(float);
+}
+
 int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);

@@ -2682,6 +2691,36 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
         ggml_metal_buffer_id bid_ids = bid_tpe;
         bid_ids.offs += ggml_metal_op_mul_mat_id_extra_tpe(op);

+        ggml_metal_buffer_id bid_amax = bid_ids;
+        bid_amax.offs += ggml_metal_op_mul_mat_id_extra_ids(op);
+
+        // src1 rescale factors, computed before the matmul
+        // ref: https://github.com/ggml-org/llama.cpp/pull/26223
+        {
+            ggml_metal_kargs_mul_mm_id_amax args = {
+                /*.ne00 =*/ ne10,
+                /*.ne01 =*/ ne11,
+                /*.ne02 =*/ ne12,
+                /*.nb01 =*/ nb11,
+                /*.nb02 =*/ nb12,
+            };
+
+            auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax_part(lib);
+
+            const size_t smem = pipeline.smem;
+
+            GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size);
+
+            ggml_metal_encoder_set_pipeline(enc, pipeline);
+            ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+            ggml_metal_encoder_set_buffer  (enc, bid_src1, 1);
+            ggml_metal_encoder_set_buffer  (enc, bid_amax, 2);
+
+            ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+
+            ggml_metal_encoder_dispatch_threadgroups(enc, N_MM_NPART_AMAX, 1, 1, 256, 1, 1);
+        }
+
         {
             ggml_metal_kargs_mul_mm_id_map0 args = {
                 ne02,
@@ -2713,7 +2752,18 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
             ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, ne02, 1, 1);
         }

-        // this barrier is always needed because the next kernel has to wait for the id maps to be computed
+        ggml_metal_op_concurrency_reset(ctx);
+
+        {
+            auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_amax(lib);
+
+            ggml_metal_encoder_set_pipeline(enc, pipeline);
+            ggml_metal_encoder_set_buffer  (enc, bid_amax, 0);
+
+            ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 32, 1, 1);
+        }
+
+        // the next kernel has to wait for the amax data
         ggml_metal_op_concurrency_reset(ctx);

         {
@@ -2745,6 +2795,7 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
             ggml_metal_encoder_set_buffer  (enc, bid_tpe,  3);
             ggml_metal_encoder_set_buffer  (enc, bid_ids,  4);
             ggml_metal_encoder_set_buffer  (enc, bid_dst,  5);
+            ggml_metal_encoder_set_buffer  (enc, bid_amax, 6);

             const size_t smem = pipeline.smem;

diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h
index 4dd8ce7af..ae72e8820 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.h
+++ b/ggml/src/ggml-metal/ggml-metal-ops.h
@@ -36,6 +36,7 @@ size_t ggml_metal_op_mul_mat_id_extra_tpe(const struct ggml_tensor * op);

 // id map [n_tokens, n_expert]
 size_t ggml_metal_op_mul_mat_id_extra_ids(const struct ggml_tensor * op);
+size_t ggml_metal_op_mul_mat_id_extra_amax(const struct ggml_tensor * op);

 // return true if we should use the FA vector kernel for this op
 bool ggml_metal_op_flash_attn_ext_use_vec(const struct ggml_tensor * op);
diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp
index 4cbec8645..4f9440f9e 100644
--- a/ggml/src/ggml-metal/ggml-metal.cpp
+++ b/ggml/src/ggml-metal/ggml-metal.cpp
@@ -226,6 +226,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
             {
                 res += ggml_metal_op_mul_mat_id_extra_tpe(tensor);
                 res += ggml_metal_op_mul_mat_id_extra_ids(tensor);
+                res += ggml_metal_op_mul_mat_id_extra_amax(tensor);
             } break;
         case GGML_OP_FLASH_ATTN_EXT:
             {
diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal
index 0a45bb1bb..71d991149 100644
--- a/ggml/src/ggml-metal/kernels/mul_mm.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mm.metal
@@ -413,6 +413,85 @@ kernel void kernel_mul_mm_id_map0(
     tpe_u32[ide] = n_all;
 }

+kernel void kernel_mul_mm_id_amax_part_f32(
+        constant ggml_metal_kargs_mul_mm_id_amax & args,
+        device   const char * src1,
+        device         char * dst,
+        threadgroup    char * shmem [[threadgroup(0)]],
+        uint  tgpig[[threadgroup_position_in_grid]],
+        ushort tiitg[[thread_index_in_threadgroup]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]],
+        ushort   ntg[[threads_per_threadgroup]]) {
+    const int nrow = args.ne01*args.ne02;
+
+    float lmax = 0.0f;
+
+    for (int ir = tgpig; ir < nrow; ir += N_MM_NPART_AMAX) {
+        const int i01 = ir % args.ne01;
+        const int i02 = ir / args.ne01;
+
+        device const float * row = (device const float *) (src1 + i02*args.nb02 + i01*args.nb01);
+
+        for (int i00 = tiitg; i00 < args.ne00; i00 += ntg) {
+            lmax = max(lmax, fabs(row[i00]));
+        }
+    }
+
+    float amax = simd_max(lmax);
+
+    threadgroup float * shared_amax = (threadgroup float *) shmem;
+
+    if (ntg > N_SIMDWIDTH) {
+        if (sgitg == 0) {
+            shared_amax[tiisg] = 0.0f;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        if (tiisg == 0) {
+            shared_amax[sgitg] = amax;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
+        amax = shared_amax[tiisg];
+        amax = simd_max(amax);
+    }
+
+    if (tiitg == 0) {
+        ((device float *) (dst + 8))[tgpig] = amax;
+    }
+}
+
+kernel void kernel_mul_mm_id_amax_f32(
+        device char * dst,
+        ushort tiitg[[thread_index_in_threadgroup]]) {
+    device const float * part = (device const float *) (dst + 8);
+
+    float amax = 0.0f;
+
+    for (int i = tiitg; i < N_MM_NPART_AMAX; i += N_SIMDWIDTH) {
+        amax = max(amax, part[i]);
+    }
+
+    amax = simd_max(amax);
+
+    if (tiitg == 0) {
+        // leave a comfortable margin below the f16 max of 65504
+        float scale = 1.0f;
+
+        // isfinite: src1 already inf/nan is not ours to fix - keep the
+        // scale at 1.0 instead of turning it into a different failure
+        if (isfinite(amax) && amax > 32768.0f) {
+            scale = exp2(ceil(log2(amax)) - 15.0f);
+        }
+
+        device float * d = (device float *) dst;
+
+        d[0] = 1.0f/scale; // exact: scale is a power of two
+        d[1] = scale;
+    }
+}
+
 typedef decltype(kernel_mul_mm_id_map0<1>) kernel_mul_mm_id_map0_t;

 template [[host_name("kernel_mul_mm_id_map0_ne20_1" )]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<1>;
@@ -433,6 +512,7 @@ kernel void kernel_mul_mm_id(
         device const char * htpe,
         device const char * hids,
         device       char * dst,
+        device const char * amax,
         threadgroup  char * shmem [[threadgroup(0)]],
         uint3  tgpig[[threadgroup_position_in_grid]],
         ushort tiitg[[thread_index_in_threadgroup]],
@@ -503,6 +583,10 @@ kernel void kernel_mul_mm_id(

     const short lb1 = (short) tiitg/NL1; // 0 .. NR1-1, this thread's row of the B tile

+    // power-of-two rescaling
+    const float s1_inv   = ((device const float *) amax)[0];
+    const float s1_scale = ((device const float *) amax)[1];
+
 #ifndef GGML_METAL_HAS_TENSOR
     S0_8x8 ma[4];
     S1_8x8 mb[2];
@@ -586,7 +670,7 @@ kernel void kernel_mul_mm_id(

                 const short ib = 4*sx + sy;

-                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
+                *(sb + 64*ib + 8*ly + lx) = loop_k + iy + i < args.ne00 ? (S1) (*((device T1 *) y + i) * (T1) s1_inv) : 0;
             }
         } else {
             const short sx = (tiitg%NL1);
@@ -599,7 +683,7 @@ kernel void kernel_mul_mm_id(

             const short ib = 4*sx + sy;

-            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)(*((device T1_2x4 *) y));
+            *(threadgroup S1_2x4 *)(sb + 64*ib + 8*ly) = (S1_2x4)((*((device T1_2x4 *) y)) * (T1) s1_inv);
         }
 #else
         // load data and store to threadgroup memory
@@ -647,7 +731,7 @@ kernel void kernel_mul_mm_id(
                 //const short lx = (tiitg/NL1)%8;
                 //const short ly = i;

-                *(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) *((device T1 *) y + i) : 0;
+                *(sb + NK*(8*sy + ly) + 8*sx + lx) = loop_k + iy + i < args.ne00 ? (S1) (*((device T1 *) y + i) * (T1) s1_inv) : 0;
             }
         } else {
             const short sx = (tiitg%NL1);
@@ -658,7 +742,7 @@ kernel void kernel_mul_mm_id(
             //const short lx = (tiitg/NL1)%8;
             //const short ly = i;

-            *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)(*((device T1_2x4 *) y));
+            *(threadgroup S1_2x4 *)(sb + NK*(8*sy + ly) + 8*sx) = (S1_2x4)((*((device T1_2x4 *) y)) * (T1) s1_inv);
         }
 #endif

@@ -749,12 +833,12 @@ kernel void kernel_mul_mm_id(

         int i = tiisg;
         for (; i < nr0/4; i += 32) {
-            *(D4 + i) = *(C4 + i);
+            *(D4 + i) = *(C4 + i) * s1_scale;
         }

         i = (4*(nr0/4)) + tiisg;
         for (; i < nr0; i += 32) {
-            *(D + i) = *(C + i);
+            *(D + i) = *(C + i) * s1_scale;
         }
     }
 }
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index bd75e2756..c7e1d7010 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -5026,9 +5026,13 @@ static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats) {
     }
 }

-static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) {
+static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats, float amax = 1.0f) {
     for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
-        if (t->type != GGML_TYPE_I32) {
+        if (t->type == GGML_TYPE_I32) {
+            continue;
+        } else if (amax != 1.0f && t->type == GGML_TYPE_F32) {
+            init_tensor_uniform(t, -amax, amax);
+        } else {
             init_tensor_uniform(t);
         }
     }
@@ -5045,9 +5049,10 @@ struct test_mul_mat_id : public test_case {
     const int64_t m;
     const int64_t n;
     const int64_t k;
+    const float amax; // magnitude of src1

     std::string vars() override {
-        return VARS_TO_STR8(type_a, type_b, n_mats, n_used, b, m, n, k);
+        return VARS_TO_STR9(type_a, type_b, n_mats, n_used, b, m, n, k, amax);
     }

     double max_nmse_err() override {
@@ -5069,9 +5074,10 @@ struct test_mul_mat_id : public test_case {

     test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
             int n_mats = 8, int n_used = 2, bool b = false,
-            int64_t m = 32, int64_t n = 32, int64_t k = 32)
+            int64_t m = 32, int64_t n = 32, int64_t k = 32,
+            float amax = 1.0f)
         : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b),
-            m(m), n(n), k(k) {
+            m(m), n(n), k(k), amax(amax) {
             GGML_ASSERT(n_used <= n_mats);
         }

@@ -5097,7 +5103,7 @@ struct test_mul_mat_id : public test_case {
     }

     void initialize_tensors(ggml_context * ctx) override {
-        init_mul_mat_id_tensors(ctx, n_mats);
+        init_mul_mat_id_tensors(ctx, n_mats, amax);
     }

     void reinit_perf_iter(ggml_context * ctx) override {
@@ -10069,6 +10075,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 4, false, 16, 10, 256));
     }

+    // test src1 f16 overflow
+    // TODO: https://github.com/ggml-org/llama.cpp/pull/26223#issuecomment-5585815365
+    //for (int n : {16, 32, 64}) {
+    //    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_K, GGML_TYPE_F32, 128, 4, false, 4096, n, 2048, 1e5f));
+    //    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8,   2, false, 512,  n, 256,  1e5f));
+    //}
+
     for (ggml_type type_a : base_types) {
         for (ggml_type type_b : {GGML_TYPE_F32 /*, GGML_TYPE_F16 */}) {
             for (int n_mats : {4, 8}) {