Sycl iq3 crash fix (#30604)

ggml/src/ggml-sycl/convert.cpp

@@ -440,7 +440,8 @@ static void dequantize_row_iq3_xxs_sycl_reorder(const void * vx, dst_t * y, cons
         sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * sycl::range<3>(1, 1, 32), sycl::range<3>(1, 1, 32)),
         [=](sycl::nd_item<3> item_ct1) {
             dequantize_block_iq3_xxs_reorder(
-                vx, y, item_ct1, nb, iq3xxs_grid, ksigns_iq2xs, kmask_iq2xs);
+                vx, y + item_ct1.get_group(2) * QK_K, item_ct1.get_group(2), item_ct1, nb,
+                iq3xxs_grid, ksigns_iq2xs, kmask_iq2xs);
         });
 }
 
@@ -472,7 +473,8 @@ static void dequantize_row_iq3_s_sycl_reorder(const void * vx, dst_t * y, const
     stream->parallel_for(
         sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * sycl::range<3>(1, 1, 32), sycl::range<3>(1, 1, 32)),
         [=](sycl::nd_item<3> item_ct1) {
-            dequantize_block_iq3_s_reorder(vx, y, item_ct1, nb, kmask_iq2xs, iq3s_grid);
+            dequantize_block_iq3_s_reorder(vx, y + item_ct1.get_group(2) * QK_K, item_ct1.get_group(2),
+                                           item_ct1, nb, kmask_iq2xs, iq3s_grid);
         });
 }
 

ggml/src/ggml-sycl/dequantize.hpp

@@ -1481,12 +1481,11 @@ static void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __res
 }
 
 template<typename dst_t>
-static void dequantize_block_iq3_xxs_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy,
+static void dequantize_block_iq3_xxs_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, const int64_t i,
                                              const sycl::nd_item<3> & item_ct1, const int64_t nb,
                                              const uint32_t * iq3xxs_grid,
                                              const uint8_t * ksigns_iq2xs,
                                              const uint8_t * kmask_iq2xs) {
-    const int64_t i   = item_ct1.get_group(2);
     const int64_t tid = item_ct1.get_local_id(2);
 #if QK_K == 256
     const int64_t il = tid / 8;
@@ -1501,7 +1500,7 @@ static void dequantize_block_iq3_xxs_reorder(const void * __restrict__ vx, dst_t
     const uint32_t aux32 = gas[0] | (gas[1] << 16);
     const float dl = (float) d[i] * (0.5f + (aux32 >> 28)) * 0.5f;
     const uint8_t signs = ksigns_iq2xs[(aux32 >> (7 * il)) & 127];
-    dst_t * y = yy + i * QK_K + 32 * ib + 8 * il;
+    dst_t * y = yy + 32 * ib + 8 * il;
 #pragma unroll
     for (int j = 0; j < 4; ++j) {
         y[j+0] = dl * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
@@ -1510,6 +1509,7 @@ static void dequantize_block_iq3_xxs_reorder(const void * __restrict__ vx, dst_t
 #else
     GGML_UNUSED(vx);
     GGML_UNUSED(yy);
+    GGML_UNUSED(i);
     GGML_UNUSED(item_ct1);
     GGML_UNUSED(nb);
     GGML_UNUSED(iq3xxs_grid);
@@ -1551,10 +1551,9 @@ dequantize_block_iq3_s(const void *__restrict__ vx, dst_t *__restrict__ yy,
 
 template <typename dst_t>
 __dpct_inline__ static void
-dequantize_block_iq3_s_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy,
+dequantize_block_iq3_s_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, const int64_t i,
                                const sycl::nd_item<3> & item_ct1, const int64_t nb,
                                const uint8_t * kmask_iq2xs, const uint32_t * iq3s_grid) {
-    const int64_t i   = item_ct1.get_group(2);
     const int64_t tid = item_ct1.get_local_id(2);
 #if QK_K == 256
     const int64_t il = tid / 8;
@@ -1574,7 +1573,7 @@ dequantize_block_iq3_s_reorder(const void * __restrict__ vx, dst_t * __restrict_
         iq3s_grid + (qs_ib[2*il+1] | ((qh[ib] << (7 - 2*il)) & 256)));
     const float dl = (float) d * (1 + 2 * ((scales[ib/2] >> (4 * (ib%2))) & 0xf));
     const uint8_t signs = signs_base[4 * ib + il];
-    dst_t * y = yy + i * QK_K + 32 * ib + 8 * il;
+    dst_t * y = yy + 32 * ib + 8 * il;
 #pragma unroll
     for (int j = 0; j < 4; ++j) {
         y[j+0] = dl * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
@@ -1583,6 +1582,7 @@ dequantize_block_iq3_s_reorder(const void * __restrict__ vx, dst_t * __restrict_
 #else
     GGML_UNUSED(vx);
     GGML_UNUSED(yy);
+    GGML_UNUSED(i);
     GGML_UNUSED(item_ct1);
     GGML_UNUSED(nb);
     GGML_UNUSED(kmask_iq2xs);

ggml/src/ggml-sycl/fused-gemm.cpp

@@ -860,7 +860,7 @@ static void fused_dequant_gemm_tile(
     }
 }
 
-template <typename S, typename block_q_t>
+template <typename S, typename block_q_t, bool reordered = false>
 static void fused_dequant_gemm_launch(const void * src0, const typename S::tsb * packed, float * dst, const int M,
                                       const int N, const int Npad, const int K, const int ldd,
                                       const int64_t groups_n, const int64_t groups_m, dpct::queue_ptr stream) {
@@ -871,8 +871,8 @@ static void fused_dequant_gemm_launch(const void * src0, const typename S::tsb *
             sycl::nd_range<2>(sycl::range<2>(groups_n, groups_m * S::WG_SIZE), sycl::range<2>(1, S::WG_SIZE)),
             [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(S::SG)]] {
                 const int n0 = item.get_group(0) * S::BN;
-                fused_dequant_gemm_tile<S, block_q_t, false>((const block_q_t *) src0, packed, dst, M, Npad, K, ldd,
-                                                      n0, n0, N, tile_a, tile_c, item);
+                fused_dequant_gemm_tile<S, block_q_t, reordered>((const block_q_t *) src0, packed, dst, M, Npad, K,
+                                                      ldd, n0, n0, N, tile_a, tile_c, item);
             });
     });
 }
@@ -974,7 +974,7 @@ template <typename F> static bool fg_visit_type(ggml_type type, bool reordered,
 }
 
 template <typename S>
-static bool fg_fused_run(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+static bool fg_fused_run(ggml_type src0_type, bool src0_reordered, const void * src0, const void * src1, ggml_type src1_type,
                          float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd, ggml_sycl_pool & pool,
                          dpct::queue_ptr stream) {
     const int64_t groups_n = (N + S::BN - 1) / S::BN;
@@ -992,21 +992,27 @@ static bool fg_fused_run(ggml_type src0_type, const void * src0, const void * sr
     }
 
     const typename S::tsb * packed = packed_b.get();
-    return fg_visit_type(src0_type, false, [&](auto tag) {
+    return fg_visit_type(src0_type, src0_reordered, [&](auto tag) {
         using block_q_t = typename decltype(tag)::type;
         if constexpr (fg_plain_ok<block_q_t>()) {
-            fused_dequant_gemm_launch<S, block_q_t>(src0, packed, dst, (int) M, (int) N, Npad, (int) K, (int) ldd,
-                                                    groups_n, groups_m, stream);
+            fused_dequant_gemm_launch<S, block_q_t, decltype(tag)::reordered>(src0, packed, dst, (int) M, (int) N,
+                                                                              Npad, (int) K, (int) ldd, groups_n,
+                                                                              groups_m, stream);
         }
     });
 }
 
-bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, bool src0_reordered, const void * src0, const void * src1, ggml_type src1_type,
                                   int32_t src1_prec, float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd,
                                   ggml_sycl_pool & pool, dpct::queue_ptr stream) {
-    // every BN columns dequantize A again, so wide N is left to the library GEMM
+    // every BN columns dequantize A again, so wide N is left to the library GEMM.
+    // src0_reordered: the persistent weight reorder may have been installed on this
+    // tensor by an earlier small-batch mul_mat; the A stage then has to decode the
+    // SoA layout. Types without an SoA stage (fg_soa<T>::supported == false)
+    // decline via fg_visit_layout and the caller falls through to the
+    // reorder-aware conversion route.
     bool plain_ok = false;
-    fg_visit_type(src0_type, false, [&](auto tag) { plain_ok = fg_plain_ok<typename decltype(tag)::type>(); });
+    fg_visit_type(src0_type, src0_reordered, [&](auto tag) { plain_ok = fg_plain_ok<typename decltype(tag)::type>(); });
     if (g_ggml_sycl_dynamic_precision == GGML_SYCL_DYNAMIC_PRECISION_F32 ||
         !ggml_sycl_xmx_gather_type_enabled(src0_type) || !plain_ok) {
         return false;
@@ -1023,7 +1029,7 @@ bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, const void * src0, const
     }
     bool launched = false;
     fg_visit_combo(combo, [&](auto s) {
-        launched = fg_fused_run<decltype(s)>(src0_type, src0, src1, src1_type, dst, M, N, K, ldd, pool, stream);
+        launched = fg_fused_run<decltype(s)>(src0_type, src0_reordered, src0, src1, src1_type, dst, M, N, K, ldd, pool, stream);
     });
     return launched;
 }

ggml/src/ggml-sycl/fused-gemm.hpp

@@ -69,7 +69,7 @@ inline bool ggml_sycl_xmx_gather_type_enabled(ggml_type src0_type) {
 // (op_params[3], [TAG_GGML_PREC]); the accumulator is f32, which meets any request.
 // q8_0 and the k-quants are not handled here, only in the grouped path below.
 // Returns false when the case is not handled (type, device, precision, or shape).
-bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, const void * src0, const void * src1, ggml_type src1_type,
+bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, bool src0_reordered, const void * src0, const void * src1, ggml_type src1_type,
                                   int32_t src1_prec, float * dst, int64_t M, int64_t N, int64_t K, int64_t ldd,
                                   ggml_sycl_pool & pool, dpct::queue_ptr stream);
 

ggml/src/ggml-sycl/getrows.cpp

@@ -301,6 +301,52 @@ static void get_rows_sycl_reorder_k(const ggml_tensor *src0, const ggml_tensor *
     GGML_UNUSED(src0);
 }
 
+// IQ3_S / IQ3_XXS in the persistent reorder layout: one work-group of 32 threads
+// dequantizes one QK_K block of one requested row, reusing the reorder-aware
+// block dequantizers that the dequantize-based mul_mat path already uses.
+template <typename F>
+static void get_rows_sycl_reorder_iq3(ggml_backend_sycl_context & ctx, const ggml_tensor * src0,
+                                      const ggml_tensor * src1, ggml_tensor * dst,
+                                      const int32_t * src1_dd, float * dst_dd,
+                                      F dequantize_block) {
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    // the persistent reorder is only installed for 2D weights
+    GGML_ASSERT(ne02 == 1 && ne03 == 1);
+    GGML_ASSERT(ne00 % QK_K == 0);
+
+    const int64_t nblk    = ne00 / QK_K;
+    const int64_t nb      = nblk * ne01; // blocks in the whole tensor
+    const int64_t ngroups = ne10 * ne11 * ne12 * nblk;
+
+    const size_t s1  = nb1 / ggml_element_size(dst);
+    const size_t s2  = nb2 / ggml_element_size(dst);
+    const size_t s3  = nb3 / ggml_element_size(dst);
+
+    const size_t s10 = nb10 / ggml_element_size(src1);
+    const size_t s11 = nb11 / ggml_element_size(src1);
+    const size_t s12 = nb12 / ggml_element_size(src1);
+
+    queue_ptr stream = ctx.stream();
+    stream->parallel_for(
+        sycl::nd_range<3>(sycl::range<3>(1, 1, ngroups * 32), sycl::range<3>(1, 1, 32)),
+        [=](sycl::nd_item<3> item_ct1) {
+            const int64_t g   = item_ct1.get_group(2);
+            const int64_t b   = g % nblk;
+            const int64_t r   = g / nblk;
+            const int64_t i10 = r % ne10;
+            const int64_t i11 = (r / ne10) % ne11;
+            const int64_t i12 = r / (ne10 * ne11);
+
+            const int64_t i01 = src1_dd[i10*s10 + i11*s11 + i12*s12];
+
+            float * y = dst_dd + i10*s1 + i11*s2 + i12*s3 + b*QK_K;
+            dequantize_block(y, i01*nblk + b, item_ct1, nb);
+        });
+
+    GGML_UNUSED(dst);
+}
+
 template <int qk, int qr, dequantize_kernel_f32_t dq>
 static void get_rows_sycl_f32(ggml_backend_sycl_context & ctx, const ggml_tensor *src0, const ggml_tensor *src1,
                               ggml_tensor *dst, const void *src0_dd,
@@ -460,7 +506,9 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type));
 
     const int32_t * src1_i32 = (const int32_t *) dst->src[1]->data;
-    // a weight shared with MUL_MAT (e.g. tied embeddings) can be in reorder layout
+    // a weight shared with MUL_MAT (e.g. tied embeddings) can be in the
+    // persistent reorder layout (incl. IQ3_S / IQ3_XXS since the SYCL IQ3
+    // reorder work)
     const bool src0_reordered = dst->src[0]->extra &&
         ((const ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.reorder;
     /* TODO: Refactor and remove duplicates */
@@ -506,6 +554,15 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_IQ3_XXS:
+            if (src0_reordered) {
+                const void * vx = dst->src[0]->data;
+                get_rows_sycl_reorder_iq3(ctx, dst->src[0], dst->src[1], dst, src1_i32, (float *) dst->data,
+                    [=](float * y, int64_t iblk, const sycl::nd_item<3> & item_ct1, int64_t nb) {
+                        dequantize_block_iq3_xxs_reorder(vx, y, iblk, item_ct1, nb,
+                                                         iq3xxs_grid, ksigns_iq2xs, kmask_iq2xs);
+                    });
+                break;
+            }
             get_rows_sycl<QK_K, 1, dequantize_iq3_xxs>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;
@@ -518,6 +575,15 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_IQ3_S:
+            if (src0_reordered) {
+                const void * vx = dst->src[0]->data;
+                get_rows_sycl_reorder_iq3(ctx, dst->src[0], dst->src[1], dst, src1_i32, (float *) dst->data,
+                    [=](float * y, int64_t iblk, const sycl::nd_item<3> & item_ct1, int64_t nb) {
+                        dequantize_block_iq3_s_reorder(vx, y, iblk, item_ct1, nb,
+                                                       kmask_iq2xs, iq3s_grid);
+                    });
+                break;
+            }
             get_rows_sycl<QK_K, 1, dequantize_iq3_s>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -3154,9 +3154,16 @@ inline void ggml_sycl_op_mul_mat_sycl(
 #endif
 
     // dequantize inside the GEMM instead of writing the f16 weights out and reading them back; src1
-    // goes in its own type, so there is no separate conversion pass
+    // goes in its own type, so there is no separate conversion pass. The A stage
+    // decodes the weight bytes in place, so an installed persistent reorder must be
+    // visible to it: pass the flag through; types without an SoA A-stage decline
+    // inside and this falls through to the reorder-aware conversion route below.
+    bool src0_reordered = false;
+    if (const auto * src0_extra = static_cast<const ggml_tensor_extra_gpu *>(src0->extra)) {
+        src0_reordered = src0_extra->optimized_feature.reorder;
+    }
     if (ggml_is_quantized(src0->type) && ggml_is_contiguous(src0) && row_diff == src0->ne[1] &&
-        ggml_sycl_fused_dequant_gemm(src0->type, src0_dd_i, src1_ddf_i, src1->type, ggml_sycl_src1_prec(dst), dst_dd_i,
+        ggml_sycl_fused_dequant_gemm(src0->type, src0_reordered, src0_dd_i, src1_ddf_i, src1->type, ggml_sycl_src1_prec(dst), dst_dd_i,
                                      row_diff, src1_ncols, ne10, ldc, ctx.pool(), stream)) {
         return;
     }

tests/test-backend-ops.cpp

@@ -2454,7 +2454,9 @@ struct test_get_rows : public test_case {
 };
 
 // GGML_OP_GET_ROWS on a weight also used by GGML_OP_MUL_MAT (tied embeddings)
-// the mul_mat runs first, so a backend that changes the weight layout must still return the right rows
+// that the preceding mul_mat may have put in the SYCL reorder layout. The
+// mul_mat runs first, so a backend that changes the weight layout in place
+// must still return the right rows from get_rows.
 struct test_get_rows_after_mul_mat : public test_case {
     const ggml_type type;
     const int n; // cols
@@ -2510,6 +2512,48 @@ struct test_get_rows_after_mul_mat : public test_case {
     }
 };
 
+// The persistent SYCL weight reorder installs on the first small-batch mul_mat; a
+// later >8-column mul_mat on the same weight goes through the fused-dequant GEMM /
+// conversion routes, which must decode the reordered layout (regression test for the
+// #29107 crash: the plain fused-dequant GEMM decoded standard layout unconditionally).
+struct test_mul_mat_reorder_then_big_batch : public test_case {
+    const ggml_type type;
+    const int n; // cols of the big batch
+    const int m; // rows
+
+    std::string vars() override {
+        return VARS_TO_STR3(type, n, m);
+    }
+
+    test_mul_mat_reorder_then_big_batch(ggml_type type = GGML_TYPE_IQ3_S, int n = 64, int m = 64)
+        : type(type), n(n), m(m) {}
+
+    double max_nmse_err() override {
+        return 5e-4; // same as test_mul_mat
+    }
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * w = ggml_new_tensor_2d(ctx, type, 1*256, m);
+        ggml_set_name(w, "w");
+
+        ggml_tensor * x1 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1*256, 1);
+        ggml_set_name(x1, "x1");
+
+        ggml_tensor * x2 = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1*256, n);
+        ggml_set_name(x2, "x2");
+
+        ggml_tensor * mm1 = ggml_mul_mat(ctx, w, x1);
+        ggml_set_name(mm1, "mm1");
+
+        ggml_tensor * mm2 = ggml_mul_mat(ctx, w, x2);
+        ggml_set_name(mm2, "mm2");
+
+        // mm1 is the first source of the final ADD, so the reorder-installing
+        // small-batch mul_mat is computed before the big-batch one
+        return ggml_add(ctx, ggml_repeat(ctx, mm1, mm2), mm2);
+    }
+};
+
 // GGML_OP_GET_ROWS_BACK
 struct test_get_rows_back : public test_case {
     const ggml_type type;
@@ -9885,9 +9929,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     }
 
     test_cases.emplace_back(new test_get_rows(GGML_TYPE_F32, 1, 8, 2, 1, 1, false));
-    for (ggml_type type : {GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K}) {
+    for (ggml_type type : {GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K,
+                          GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S}) {
         test_cases.emplace_back(new test_get_rows_after_mul_mat(type));
     }
+    for (ggml_type type : {GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ3_S}) {
+        test_cases.emplace_back(new test_mul_mat_reorder_then_big_batch(type));
+    }
     for (ggml_type type : all_types) {
         for (int b : {1, 7}) {
             for (bool v : {false, true}) {