From 9faa9ee7b0e54828a21149e2e02bdd04123a696b Mon Sep 17 00:00:00 2001 From: Neo Zhang Date: Fri, 7 Aug 2026 13:22:23 +0800 Subject: [PATCH] sycl : Support DSv4 OPs: LIGHTNING_INDEXER,DSV4_HC_COMB,DSV4_HC_POST,DSV4_HC_PRE (llama/26568) * support DSv4 OPs: LIGHTNING_INDEXER,DSV4_HC_COMB,DSV4_HC_POST,DSV4_HC_PREwq * update ops.md * fix format issue --- ggml/src/ggml-sycl/dsv4-hc.cpp | 280 +++++++++++++++++++++++ ggml/src/ggml-sycl/dsv4-hc.hpp | 10 + ggml/src/ggml-sycl/ggml-sycl.cpp | 35 +++ ggml/src/ggml-sycl/lightning-indexer.cpp | 197 ++++++++++++++++ ggml/src/ggml-sycl/lightning-indexer.hpp | 8 + 5 files changed, 530 insertions(+) create mode 100644 ggml/src/ggml-sycl/dsv4-hc.cpp create mode 100644 ggml/src/ggml-sycl/dsv4-hc.hpp create mode 100644 ggml/src/ggml-sycl/lightning-indexer.cpp create mode 100644 ggml/src/ggml-sycl/lightning-indexer.hpp diff --git a/ggml/src/ggml-sycl/dsv4-hc.cpp b/ggml/src/ggml-sycl/dsv4-hc.cpp new file mode 100644 index 000000000..bb66e8c1b --- /dev/null +++ b/ggml/src/ggml-sycl/dsv4-hc.cpp @@ -0,0 +1,280 @@ +#include "ggml-impl.h" +#include "dsv4-hc.hpp" + +#include + +static constexpr int DSV4_HC = 4; + +static void dsv4_hc_pre_f32_sycl( + const float * x, const float * weights, float * dst, + int64_t n_embd, int64_t hc, int64_t n_tokens, + int64_t sx0, int64_t sx1, int64_t sx2, + int64_t sw0, int64_t sw1, + int64_t sd0, int64_t sd1, + queue_ptr stream) { + const int64_t nr = n_embd * n_tokens; + const int64_t block_size = 256; + const int64_t num_blocks = (nr + block_size - 1) / block_size; + + stream->parallel_for( + sycl::nd_range<1>(sycl::range<1>(num_blocks * block_size), sycl::range<1>(block_size)), + [=](sycl::nd_item<1> item) { + const int64_t ir = item.get_global_id(0); + if (ir >= nr) { + return; + } + + const int64_t i0 = ir % n_embd; + const int64_t it = ir / n_embd; + + float sum = x[i0*sx0 + it*sx2] * weights[it*sw1]; + for (int64_t ih = 1; ih < hc; ++ih) { + const float xv = x[i0*sx0 + ih*sx1 + it*sx2]; + const float wv = weights[ih*sw0 + it*sw1]; + sum += xv * wv; + } + + dst[i0*sd0 + it*sd1] = sum; + }); +} + +static void dsv4_hc_comb_norm_cols(float * comb, float eps) { + for (int idst = 0; idst < DSV4_HC; ++idst) { + float sum = eps; + for (int isrc = 0; isrc < DSV4_HC; ++isrc) { + sum += comb[idst + DSV4_HC*isrc]; + } + + const float inv_sum = 1.0f / sum; + for (int isrc = 0; isrc < DSV4_HC; ++isrc) { + comb[idst + DSV4_HC*isrc] *= inv_sum; + } + } +} + +static void dsv4_hc_comb_norm_rows(float * comb, float eps) { + for (int isrc = 0; isrc < DSV4_HC; ++isrc) { + float sum = eps; + for (int idst = 0; idst < DSV4_HC; ++idst) { + sum += comb[idst + DSV4_HC*isrc]; + } + + const float inv_sum = 1.0f / sum; + for (int idst = 0; idst < DSV4_HC; ++idst) { + comb[idst + DSV4_HC*isrc] *= inv_sum; + } + } +} + +static void dsv4_hc_comb_f32_sycl( + const float * mixes, + const float * scale, + const float * base, + float * dst, + int64_t n_tokens, + int64_t sm0, + int64_t sm1, + int64_t ss0, + int64_t sb0, + int64_t sd0, + int64_t sd1, + int64_t sd2, + float eps, + int32_t n_iter, + queue_ptr stream) { + constexpr int comb_offset = 2*DSV4_HC; + + const int64_t block_size = 256; + const int64_t num_blocks = (n_tokens + block_size - 1) / block_size; + + stream->parallel_for( + sycl::nd_range<1>(sycl::range<1>(num_blocks * block_size), sycl::range<1>(block_size)), + [=](sycl::nd_item<1> item_ct1) { + const int64_t it = item_ct1.get_global_id(0); + + if (it >= n_tokens) { + return; + } + + const float scale_comb = scale[2*ss0]; + float comb[DSV4_HC*DSV4_HC]; + + for (int isrc = 0; isrc < DSV4_HC; ++isrc) { + float max = -INFINITY; + for (int idst = 0; idst < DSV4_HC; ++idst) { + const int idx = idst + DSV4_HC*isrc; + const float v = mixes[(comb_offset + idx)*sm0 + it*sm1] * scale_comb + base[(comb_offset + idx)*sb0]; + comb[idx] = v; + max = fmaxf(max, v); + } + + float sum = 0.0f; + for (int idst = 0; idst < DSV4_HC; ++idst) { + const int idx = idst + DSV4_HC*isrc; + const float v = expf(comb[idx] - max); + comb[idx] = v; + sum += v; + } + + const float inv_sum = 1.0f / sum; + for (int idst = 0; idst < DSV4_HC; ++idst) { + const int idx = idst + DSV4_HC*isrc; + comb[idx] = comb[idx] * inv_sum + eps; + } + } + + dsv4_hc_comb_norm_cols(comb, eps); + for (int32_t i = 1; i < n_iter; ++i) { + dsv4_hc_comb_norm_rows(comb, eps); + dsv4_hc_comb_norm_cols(comb, eps); + } + + for (int isrc = 0; isrc < DSV4_HC; ++isrc) { + for (int idst = 0; idst < DSV4_HC; ++idst) { + const int idx = idst + DSV4_HC*isrc; + dst[idst*sd0 + isrc*sd1 + it*sd2] = comb[idx]; + } + } + }); +} + +static void dsv4_hc_post_f32_sycl( + const float * x, const float * residual, const float * post, const float * comb, float * dst, + int64_t n_embd, int64_t hc, int64_t n_tokens, + int64_t sx0, int64_t sx1, + int64_t sr0, int64_t sr1, int64_t sr2, + int64_t sp0, int64_t sp1, + int64_t sc0, int64_t sc1, int64_t sc2, + int64_t sd0, int64_t sd1, int64_t sd2, + queue_ptr stream) { + const int64_t nr = n_embd * hc * n_tokens; + const int64_t block_size = 256; + const int64_t num_blocks = (nr + block_size - 1) / block_size; + + stream->parallel_for( + sycl::nd_range<1>(sycl::range<1>(num_blocks * block_size), sycl::range<1>(block_size)), + [=](sycl::nd_item<1> item) { + const int64_t ir = item.get_global_id(0); + if (ir >= nr) { + return; + } + + const int64_t i0 = ir % n_embd; + const int64_t idst = (ir / n_embd) % hc; + const int64_t it = ir / (n_embd * hc); + + float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1]; + for (int64_t isrc = 0; isrc < hc; ++isrc) { + sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2]; + } + + dst[i0*sd0 + idst*sd1 + it*sd2] = sum; + }); +} + +void ggml_sycl_op_dsv4_hc_pre(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2); + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * weights = dst->src[1]; + + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_TENSOR_LOCALS(size_t, nbx, x, nb); + GGML_TENSOR_LOCALS(size_t, nbw, weights, nb); + GGML_TENSOR_LOCALS(size_t, nbd, dst, nb); + + const int64_t n_embd = x->ne[0]; + const int64_t hc = x->ne[1]; + const int64_t n_tokens = x->ne[2]; + + queue_ptr stream = ctx.stream(); + + dsv4_hc_pre_f32_sycl( + (const float *) x->data, (const float *) weights->data, (float *) dst->data, + n_embd, hc, n_tokens, + nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float), + nbw0 / sizeof(float), nbw1 / sizeof(float), + nbd0 / sizeof(float), nbd1 / sizeof(float), + stream); +} + +void ggml_sycl_op_dsv4_hc_comb(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3); + + const ggml_tensor * mixes = dst->src[0]; + const ggml_tensor * scale = dst->src[1]; + const ggml_tensor * base = dst->src[2]; + + GGML_ASSERT(mixes->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + constexpr int64_t hc_mix_dim = (2 + DSV4_HC)*DSV4_HC; + + GGML_ASSERT(mixes->ne[0] == hc_mix_dim); + GGML_ASSERT(dst->ne[0] == DSV4_HC); + GGML_ASSERT(dst->ne[1] == DSV4_HC); + GGML_ASSERT(dst->ne[2] == mixes->ne[1]); + GGML_ASSERT(scale->ne[0] >= 3); + GGML_ASSERT(base->ne[0] == hc_mix_dim); + + GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb); + GGML_TENSOR_LOCALS(size_t, nbs, scale, nb); + GGML_TENSOR_LOCALS(size_t, nbb, base, nb); + GGML_TENSOR_LOCALS(size_t, nbd, dst, nb); + + const int64_t n_tokens = mixes->ne[1]; + const float eps = ggml_get_op_params_f32(dst, 0); + const int32_t n_iter = ggml_get_op_params_i32(dst, 1); + + queue_ptr stream = ctx.stream(); + + dsv4_hc_comb_f32_sycl( + (const float *) mixes->data, (const float *) scale->data, (const float *) base->data, (float *) dst->data, + n_tokens, + nbm0 / sizeof(float), nbm1 / sizeof(float), + nbs0 / sizeof(float), + nbb0 / sizeof(float), + nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float), + eps, n_iter, stream); +} + +void ggml_sycl_op_dsv4_hc_post(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/4); + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * residual = dst->src[1]; + const ggml_tensor * post = dst->src[2]; + const ggml_tensor * comb = dst->src[3]; + + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(residual->type == GGML_TYPE_F32); + GGML_ASSERT(post->type == GGML_TYPE_F32); + GGML_ASSERT(comb->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_TENSOR_LOCALS(size_t, nbx, x, nb); + GGML_TENSOR_LOCALS(size_t, nbr, residual, nb); + GGML_TENSOR_LOCALS(size_t, nbp, post, nb); + GGML_TENSOR_LOCALS(size_t, nbc, comb, nb); + GGML_TENSOR_LOCALS(size_t, nbd, dst, nb); + + const int64_t n_embd = x->ne[0]; + const int64_t n_tokens = x->ne[1]; + const int64_t hc = residual->ne[1]; + + queue_ptr stream = ctx.stream(); + + dsv4_hc_post_f32_sycl( + (const float *) x->data, (const float *) residual->data, + (const float *) post->data, (const float *) comb->data, (float *) dst->data, + n_embd, hc, n_tokens, + nbx0 / sizeof(float), nbx1 / sizeof(float), + nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float), + nbp0 / sizeof(float), nbp1 / sizeof(float), + nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float), + nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float), + stream); +} diff --git a/ggml/src/ggml-sycl/dsv4-hc.hpp b/ggml/src/ggml-sycl/dsv4-hc.hpp new file mode 100644 index 000000000..330518d8a --- /dev/null +++ b/ggml/src/ggml-sycl/dsv4-hc.hpp @@ -0,0 +1,10 @@ +#ifndef GGML_SYCL_DSV4_HC_HPP +#define GGML_SYCL_DSV4_HC_HPP + +#include "common.hpp" + +void ggml_sycl_op_dsv4_hc_pre(ggml_backend_sycl_context & ctx, ggml_tensor * dst); +void ggml_sycl_op_dsv4_hc_comb(ggml_backend_sycl_context & ctx, ggml_tensor * dst); +void ggml_sycl_op_dsv4_hc_post(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + +#endif // GGML_SYCL_DSV4_HC_HPP diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index ce92d438d..18d58782e 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -62,6 +62,8 @@ #include "ggml-sycl/repeat_back.hpp" #include "ggml-sycl/set_rows.hpp" #include "ggml-sycl/set.hpp" +#include "ggml-sycl/dsv4-hc.hpp" +#include "ggml-sycl/lightning-indexer.hpp" #include "ggml-sycl/conv2d.hpp" #include "ggml-sycl/conv2d-dw.hpp" #include "ggml-sycl/conv2d-transpose.hpp" @@ -4942,6 +4944,18 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg case GGML_OP_SET_ROWS: ggml_sycl_op_set_rows(ctx, dst); break; + case GGML_OP_DSV4_HC_PRE: + ggml_sycl_op_dsv4_hc_pre(ctx, dst); + break; + case GGML_OP_DSV4_HC_COMB: + ggml_sycl_op_dsv4_hc_comb(ctx, dst); + break; + case GGML_OP_DSV4_HC_POST: + ggml_sycl_op_dsv4_hc_post(ctx, dst); + break; + case GGML_OP_LIGHTNING_INDEXER: + ggml_sycl_op_lightning_indexer(ctx, dst); + break; case GGML_OP_DUP: ggml_sycl_dup(ctx, dst); break; @@ -5801,6 +5815,27 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons return res; } break; + case GGML_OP_DSV4_HC_PRE: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32; + case GGML_OP_DSV4_HC_COMB: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; + case GGML_OP_DSV4_HC_POST: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32; + case GGML_OP_LIGHTNING_INDEXER: + return op->src[0]->type == GGML_TYPE_F32 && + (op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32 || + op->src[1]->type == GGML_TYPE_BF16 || op->src[1]->type == GGML_TYPE_Q8_0 || + op->src[1]->type == GGML_TYPE_Q5_1 || op->src[1]->type == GGML_TYPE_Q5_0 || + op->src[1]->type == GGML_TYPE_Q4_1 || op->src[1]->type == GGML_TYPE_Q4_0 || + op->src[1]->type == GGML_TYPE_IQ4_NL) && + op->src[2]->type == GGML_TYPE_F32 && + op->src[3]->type == GGML_TYPE_F16 && + op->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == WARP_SIZE * 8; case GGML_OP_CPY: { ggml_type src0_type = op->src[0]->type; diff --git a/ggml/src/ggml-sycl/lightning-indexer.cpp b/ggml/src/ggml-sycl/lightning-indexer.cpp new file mode 100644 index 000000000..823b713ca --- /dev/null +++ b/ggml/src/ggml-sycl/lightning-indexer.cpp @@ -0,0 +1,197 @@ +#include "lightning-indexer.hpp" +#include "dequantize.hpp" + +static void lightning_indexer_f32_sycl( + const char * q, const char * k, const char * w, const char * m, float * dst, + int64_t n_embd, int64_t n_head, int64_t n_batch, int64_t n_stream, int64_t n_kv, + int64_t nem3, + int64_t nbq1, int64_t nbq2, int64_t nbq3, + int64_t nbk2, int64_t nbk3, + int64_t nbw1, int64_t nbw3, + int64_t nbm1, int64_t nbm3, + int64_t nb1, int64_t nb3, + ggml_type k_type, + queue_ptr stream) { + + constexpr int64_t LANES = WARP_SIZE; + constexpr int64_t ELEMS_PER_LANE = 8; + constexpr int64_t ROWS_PER_BLOCK = 4; + constexpr int64_t BLOCK_SIZE = ROWS_PER_BLOCK * LANES; + + const int64_t n_rows = n_batch * n_stream * n_kv; + const int64_t n_blocks = (n_rows + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK; + + stream->parallel_for( + sycl::nd_range<1>( + sycl::range<1>(n_blocks * BLOCK_SIZE), + sycl::range<1>(BLOCK_SIZE)), + [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + const int64_t ir = item.get_global_id(0); + const int64_t lane = ir % LANES; + const int64_t row = ir / LANES; + if (row >= n_rows) { + return; + } + + const int64_t i_bs = row / n_kv; + const int64_t i_kv = row % n_kv; + const int64_t i_batch = i_bs / n_stream; + const int64_t i_stream = i_bs % n_stream; + + // load K row slice into registers (row is contiguous, nbk0 == type size) + const char * k_base = k + i_kv*nbk2 + i_stream*nbk3; + float k_local[ELEMS_PER_LANE]; + if (k_type == GGML_TYPE_F16) { + const sycl::half * k_row = (const sycl::half *) k_base; +#pragma unroll + for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) { + k_local[j] = static_cast(k_row[lane*ELEMS_PER_LANE + j]); + } + } else if (k_type == GGML_TYPE_F32) { + const float * k_row = (const float *) k_base; +#pragma unroll + for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) { + k_local[j] = k_row[lane*ELEMS_PER_LANE + j]; + } + } else { + const int64_t lane_base = lane * ELEMS_PER_LANE; + switch (k_type) { + case GGML_TYPE_BF16: { + const sycl::ext::oneapi::bfloat16 * k_row = (const sycl::ext::oneapi::bfloat16 *) k_base; +#pragma unroll + for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) { + k_local[j] = static_cast(k_row[lane_base + j]); + } + } break; + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: { +#pragma unroll + for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) { + const int64_t idx = lane_base + j; + const int64_t ib = idx / QK4_0; + const int iqs = idx % (QK4_0/2); + dfloat2 kv; + if (k_type == GGML_TYPE_Q4_0) { + dequantize_q4_0(k_base, ib, iqs, kv); + } else if (k_type == GGML_TYPE_Q4_1) { + dequantize_q4_1(k_base, ib, iqs, kv); + } else if (k_type == GGML_TYPE_Q5_0) { + dequantize_q5_0(k_base, ib, iqs, kv); + } else { + dequantize_q5_1(k_base, ib, iqs, kv); + } + k_local[j] = (idx % QK4_0) < (QK4_0/2) ? static_cast(kv.x()) : static_cast(kv.y()); + } + } break; + case GGML_TYPE_Q8_0: { +#pragma unroll + for (int64_t pair = 0; pair < ELEMS_PER_LANE / 2; ++pair) { + const int64_t elem0 = lane_base + 2 * pair; + dfloat2 kv; + dequantize_q8_0(k_base, elem0 / QK8_0, elem0 % QK8_0, kv); + k_local[2 * pair + 0] = static_cast(kv.x()); + k_local[2 * pair + 1] = static_cast(kv.y()); + } + } break; + case GGML_TYPE_IQ4_NL: { +#pragma unroll + for (int64_t pair = 0; pair < ELEMS_PER_LANE / 2; ++pair) { + const int64_t elem0 = lane_base + 2 * pair; + dfloat2 kv; + dequantize_iq4_nl(k_base, elem0 / QK4_NL, elem0 % QK4_NL, kv); + k_local[2 * pair + 0] = static_cast(kv.x()); + k_local[2 * pair + 1] = static_cast(kv.y()); + } + } break; + default: +#pragma unroll + for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) { + k_local[j] = 0.0f; + } + break; + } + } + + const char * q_base = q + i_batch*nbq2 + i_stream*nbq3; + const float * w_base = (const float *) (w + i_batch*nbw1 + i_stream*nbw3); + + float score = 0.0f; + for (int64_t h = 0; h < n_head; ++h) { + const float * q_row = (const float *) (q_base + h*nbq1); + float dot = 0.0f; +#pragma unroll + for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) { + const int64_t i = lane*ELEMS_PER_LANE + j; + if (i < n_embd) { + dot += q_row[i] * k_local[j]; + } + } + dot = sycl::reduce_over_group(item.get_sub_group(), dot, sycl::plus()); + if (lane == 0) { + score += sycl::max(dot, 0.0f) * w_base[h]; + } + } + + if (lane == 0) { + const sycl::half * m_base = (const sycl::half *) (m + i_batch*nbm1 + (i_stream % nem3)*nbm3); + // flat-index store: storing through a strided base pointer + // hangs/misroutes writes on this stack when n_batch*n_stream > 1 + const int64_t dst_idx = i_kv + i_batch*(nb1/sizeof(float)) + i_stream*(nb3/sizeof(float)); + dst[dst_idx] = score + static_cast(m_base[i_kv]); + } + }); +} + +void ggml_sycl_op_lightning_indexer(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/4); + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * w = dst->src[2]; // weights + const ggml_tensor * m = dst->src[3]; // mask + + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT( q->type == GGML_TYPE_F32); + GGML_ASSERT( w->type == GGML_TYPE_F32); + GGML_ASSERT( m->type == GGML_TYPE_F16); + GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_BF16 || + k->type == GGML_TYPE_Q8_0 || k->type == GGML_TYPE_Q5_1 || k->type == GGML_TYPE_Q5_0 || + k->type == GGML_TYPE_Q4_1 || k->type == GGML_TYPE_Q4_0 || k->type == GGML_TYPE_IQ4_NL); + + GGML_TENSOR_LOCALS(int64_t, neq, q, ne); + GGML_TENSOR_LOCALS(size_t, nbq, q, nb); + GGML_TENSOR_LOCALS(int64_t, nek, k, ne); + GGML_TENSOR_LOCALS(size_t, nbk, k, nb); + GGML_TENSOR_LOCALS(size_t, nbw, w, nb); + GGML_TENSOR_LOCALS(int64_t, nem, m, ne); + GGML_TENSOR_LOCALS(size_t, nbm, m, nb); + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne); + GGML_TENSOR_LOCALS(size_t, nb, dst, nb); + + // input rows must be contiguous + GGML_ASSERT(nbq0 == ggml_type_size(q->type)); + GGML_ASSERT(nbk0 == ggml_type_size(k->type)); + GGML_ASSERT(nbm0 == ggml_type_size(m->type)); + GGML_ASSERT(nb0 == ggml_type_size(dst->type)); + + const int64_t n_embd = neq0; + const int64_t n_head = neq1; + const int64_t n_batch = neq2; + const int64_t n_stream = neq3; + const int64_t n_kv = nek2; + + GGML_ASSERT(n_embd == WARP_SIZE * 8); + + lightning_indexer_f32_sycl( + (const char *) q->data, (const char *) k->data, + (const char *) w->data, (const char *) m->data, (float *) dst->data, + n_embd, n_head, n_batch, n_stream, n_kv, nem3, + nbq1, nbq2, nbq3, + nbk2, nbk3, + nbw1, nbw3, + nbm1, nbm3, + nb1, nb3, + k->type, + ctx.stream()); +} diff --git a/ggml/src/ggml-sycl/lightning-indexer.hpp b/ggml/src/ggml-sycl/lightning-indexer.hpp new file mode 100644 index 000000000..0b88c418e --- /dev/null +++ b/ggml/src/ggml-sycl/lightning-indexer.hpp @@ -0,0 +1,8 @@ +#ifndef GGML_SYCL_LIGHTNING_INDEXER_HPP +#define GGML_SYCL_LIGHTNING_INDEXER_HPP + +#include "common.hpp" + +void ggml_sycl_op_lightning_indexer(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + +#endif // GGML_SYCL_LIGHTNING_INDEXER_HPP