metal: OUT_PROD multiplies in float staging; a training gradient overflowed half

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -812,6 +812,41 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext(ggml_
     return res;
 }
 
+// F32 x F32 with FLOAT staging (kernel_mul_mm_f32_f32_fp32), for training's internal products:
+// a gradient overflows half's range where an activation does not. Simdgroup kernel only: the
+// caller uses it when the tensor API is off.
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_fp32(ggml_metal_library_t lib, const ggml_tensor * op) {
+    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32);
+    GGML_ASSERT(!ggml_metal_device_get_props(ggml_metal_library_get_device(lib))->has_tensor);
+    char name[256];
+    const bool bc_inp = op->src[0]->ne[0] % 32 != 0;
+    const bool bc_out = op->ne[0] % 64 != 0 || op->ne[1] % 32 != 0;
+    GGML_ASSERT(op->src[1]->ne[2] <= INT16_MAX && op->src[1]->ne[3] <= INT16_MAX);
+    const int16_t ne12 = (int16_t) op->src[1]->ne[2];
+    const int16_t ne13 = (int16_t) op->src[1]->ne[3];
+    const int16_t r2   = (int16_t) (ne12 / op->src[0]->ne[2]);
+    const int16_t r3   = (int16_t) (ne13 / op->src[0]->ne[3]);
+    const char * base = "kernel_mul_mm_f32_f32_fp32";
+    snprintf(name, 256, "%s_bci=%d_bco=%d_ne12=%d_ne13=%d_r2=%d_r3=%d", base, bc_inp, bc_out, ne12, ne13, r2, r3);
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        ggml_metal_cv_t cv = ggml_metal_cv_init();
+        ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0);
+        ggml_metal_cv_set_bool(cv, bc_out, FC_MUL_MM + 1);
+        ggml_metal_cv_set_int16(cv, ne12,  FC_MUL_MM + 2);
+        ggml_metal_cv_set_int16(cv, ne13,  FC_MUL_MM + 3);
+        ggml_metal_cv_set_int16(cv, r2,    FC_MUL_MM + 4);
+        ggml_metal_cv_set_int16(cv, r3,    FC_MUL_MM + 5);
+        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
+        ggml_metal_cv_free(cv);
+    }
+    res.nr0  = 64;
+    res.nr1  = 32;
+    res.smem = 64*32*sizeof(float) + 32*32*sizeof(float); // sa + sb; the bc_out tile (8192) fits
+    res.nsg  = N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y;
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_metal_library_t lib, const ggml_tensor * op) {
     char base[256];
     char name[256];

ggml/src/ggml-metal/ggml-metal-device.h

@@ -140,6 +140,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_del
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext        (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg);
 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_mm_fp32       (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         (ggml_metal_library_t lib, const struct ggml_tensor * op);

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -2470,7 +2470,7 @@ int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) {
 
 // The F32/quantized simdgroup matmul encode, over an explicit op (its src[0], src[1] and
 // itself may be synthetic views: OUT_PROD composes it over scratch buffers).
-static void ggml_metal_encode_mul_mm(ggml_metal_op_t ctx, const ggml_tensor * op) {
+static void ggml_metal_encode_mul_mm(ggml_metal_op_t ctx, const ggml_tensor * op, bool fp32_staging = false) {
     ggml_metal_library_t lib = ctx->lib;
     ggml_metal_encoder_t enc = ctx->enc;
 
@@ -2484,7 +2484,8 @@ static void ggml_metal_encode_mul_mm(ggml_metal_op_t ctx, const ggml_tensor * op
     const int16_t r2 = ne12/ne02;
     const int16_t r3 = ne13/ne03;
 
-        auto pipeline = ggml_metal_library_get_pipeline_mul_mm(lib, op);
+        auto pipeline = fp32_staging ? ggml_metal_library_get_pipeline_mul_mm_fp32(lib, op)
+                                     : ggml_metal_library_get_pipeline_mul_mm(lib, op);
 
         ggml_metal_kargs_mul_mm args = {
             /*.ne00 =*/ ne00,
@@ -4185,7 +4186,11 @@ int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) {
         p.op = GGML_OP_MUL_MAT;
         p.src[0] = &x;
         p.src[1] = &y;
-        ggml_metal_encode_mul_mm(ctx, &p);
+        // FLOAT staging when the simdgroup kernel runs: X and Y carry gradients, which pass
+        // half's range (65504) where activations do not; half staging turned them into Inf
+        // and cost every product ~1e-3 relative precision (2026-10-06, Qwen3.5 0.8B backward)
+        const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(ctx->lib))->has_tensor;
+        ggml_metal_encode_mul_mm(ctx, &p, !has_tensor);
         ggml_metal_op_concurrency_reset(ctx);
 
         if (!first) {

ggml/src/ggml-metal/kernels/mul_mm.metal

@@ -157,8 +157,11 @@ kernel void kernel_mul_mm(
         ushort tiitg[[thread_index_in_threadgroup]],
         ushort sgitg[[simdgroup_index_in_threadgroup]]) {
 
+    // sa holds NR0 x NK = 64 x 32 staged values of S0; sb follows it. Sized by S0 so a
+    // float-staged instantiation (training: gradients exceed half's range) lays out right;
+    // for every half instantiation this is the 4096 it always was.
     threadgroup S0 * sa = (threadgroup S0 *)(shmem);
-    threadgroup S1 * sb = (threadgroup S1 *)(shmem + 4096);
+    threadgroup S1 * sb = (threadgroup S1 *)(shmem + 64*32*sizeof(S0));
 
     constexpr int NR0 = 64;
     constexpr int NR1 = 32;
@@ -770,6 +773,16 @@ template [[host_name("kernel_mul_mm_iq4_nl_f32")]]  kernel mul_mm_t kernel_mul_m
 template [[host_name("kernel_mul_mm_iq4_xs_f32")]]  kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_iq4_xs,  QK_NL, dequantize_iq4_xs,  float,  float4x4,  float, float2x4>;
 template [[host_name("kernel_mul_mm_tq2_0_f32")]]   kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_tq2_0,   QK_NL, dequantize_tq2_0,   float,  float4x4,  float, float2x4>;
 
+// FLOAT staging, for training: a gradient passes half's range (65504) where an activation
+// does not (Qwen3.5 0.8B LoRA backward, 2026-10-06: |grad| 131273, 95 values over, the
+// staged Inf turned 32k outputs of one OUT_PROD non-finite), and half staging also cost every
+// OUT_PROD ~1e-3 relative precision. Threadgroup memory: 64*32*4 + 32*32*4 = 12288 bytes.
+// The simdgroup kernel only (the tensor-API kernel has its own staging); the host asks for
+// this variant only when the tensor API is off.
+#ifndef GGML_METAL_HAS_TENSOR
+typedef decltype(kernel_mul_mm<float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>) mul_mm_fp32_t;
+template [[host_name("kernel_mul_mm_f32_f32_fp32")]] kernel mul_mm_fp32_t kernel_mul_mm<float, float4x4, simdgroup_float8x8, float, float2x4, simdgroup_float8x8, float4x4, 1, dequantize_f32, float, float4x4, float, float2x4>;
+#endif
 template [[host_name("kernel_mul_mm_f32_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_f16_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   half, half2x4>;
 template [[host_name("kernel_mul_mm_q1_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  half, half2x4>;

tests/test-backend-ops.cpp

@@ -5031,9 +5031,23 @@ struct test_out_prod : public test_case {
     // address against out's buffer passed them all (Metal's F32 transposed path did, and read
     // stale scratch in every LoRA backward on Apple silicon until it was based on a)
     const bool a_own_buffer;
+    // b's values span [-b_range, b_range]: a training gradient passes half's range (65504)
+    // where an activation does not, and a kernel that staged b through half turned those
+    // into Inf and the products into NaN (Metal, 2026-10-06)
+    const float b_range;
 
     std::string vars() override {
-        return VARS_TO_STR9(type_a, type_b, m, n, k, bs, nr, trans_b, a_own_buffer);
+        return VARS_TO_STR10(type_a, type_b, m, n, k, bs, nr, trans_b, a_own_buffer, b_range);
+    }
+
+    void initialize_tensors(ggml_context * ctx) override {
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->view_src) {
+                continue; // a view's values are its source's
+            }
+            const bool is_b = std::string(ggml_get_name(t)).rfind("b", 0) == 0;
+            init_tensor_uniform(t, is_b ? -b_range : -1.0f, is_b ? b_range : 1.0f);
+        }
     }
 
     double max_nmse_err() override {
@@ -5046,8 +5060,8 @@ struct test_out_prod : public test_case {
             int64_t m = 32, int64_t n = 32, int64_t k = 32,
             std::array<int64_t, 2> bs = {10, 10},
             std::array<int64_t, 2> nr = {2, 2},
-            bool trans_b = false, bool a_own_buffer = false)
-        : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), trans_b(trans_b), a_own_buffer(a_own_buffer) {}
+            bool trans_b = false, bool a_own_buffer = false, float b_range = 1.0f)
+        : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), trans_b(trans_b), a_own_buffer(a_own_buffer), b_range(b_range) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         return build_graph(ctx, nullptr);
@@ -5060,6 +5074,7 @@ struct test_out_prod : public test_case {
         ggml_tensor * b;
         if (trans_b) {
             b = ggml_new_tensor_4d(ctx, type_b, k, n, bs[0]*nr[0], bs[1]*nr[1]);
+            ggml_set_name(b, "b_base");
             b = ggml_transpose(ctx, b);
         } else {
             b = ggml_new_tensor_4d(ctx, type_b, n, k, bs[0]*nr[0], bs[1]*nr[1]);
@@ -9700,6 +9715,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_out_prod(GGML_TYPE_Q4_K, GGML_TYPE_F32,
                                                   256, 8, 600000, {1, 1}, {1, 1}, trans_b));
     }
+    // gradient-sized b (past half's 65504), both layouts, F32 and q8_0 weights: the values a
+    // LoRA backward actually hands OUT_PROD (|grad| 131273 measured on the 0.8B)
+    for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_Q8_0}) {
+        for (bool trans_b : {false, true}) {
+            test_cases.emplace_back(new test_out_prod(type_a, GGML_TYPE_F32, 256, 64, 64, {1, 1}, {1, 1}, trans_b, false, 1.5e5f));
+        }
+    }
     // a in its own buffer (a LoRA weight in the adapter's buffer), F32 and quantized, both
     // gradient layouts, at LoRA's shape (rank 8 against a 1024-wide activation)
     for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_Q8_0}) {