diff --git a/ggml/src/ggml-sycl/dsv4-hc.cpp b/ggml/src/ggml-sycl/dsv4-hc.cpp index bb66e8c1b4..337f4af455 100644 --- a/ggml/src/ggml-sycl/dsv4-hc.cpp +++ b/ggml/src/ggml-sycl/dsv4-hc.cpp @@ -2,22 +2,30 @@ #include "dsv4-hc.hpp" #include +#include static constexpr int DSV4_HC = 4; +// tunable: one work-item per (embedding element, token) +static constexpr int dsv4_hc_pre_block_size = 256; + +// gated: the weight is a per-element gate [n_embd, hc, n_tokens] passed through a sigmoid. +// otherwise it is one weight per (stream, token). +template 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 sw0, int64_t sw1, int64_t sw2, int64_t sd0, int64_t sd1, + float scale, 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; + const int64_t num_blocks = (nr + dsv4_hc_pre_block_size - 1) / dsv4_hc_pre_block_size; stream->parallel_for( - sycl::nd_range<1>(sycl::range<1>(num_blocks * block_size), sycl::range<1>(block_size)), + sycl::nd_range<1>(sycl::range<1>(num_blocks * dsv4_hc_pre_block_size), + sycl::range<1>(dsv4_hc_pre_block_size)), [=](sycl::nd_item<1> item) { const int64_t ir = item.get_global_id(0); if (ir >= nr) { @@ -27,14 +35,20 @@ static void dsv4_hc_pre_f32_sycl( 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) { + float sum = 0.0f; + for (int64_t ih = 0; ih < hc; ++ih) { const float xv = x[i0*sx0 + ih*sx1 + it*sx2]; - const float wv = weights[ih*sw0 + it*sw1]; + float wv; + if constexpr (gated) { + const float gv = weights[i0*sw0 + ih*sw1 + it*sw2]; + wv = 1.0f / (1.0f + sycl::exp(-gv)); + } else { + wv = weights[ih*sw0 + it*sw1]; + } sum += xv * wv; } - dst[i0*sd0 + it*sd1] = sum; + dst[i0*sd0 + it*sd1] = scale * sum; }); } @@ -138,6 +152,12 @@ static void dsv4_hc_comb_f32_sycl( }); } +// tunable: one work-item per (embedding element, stream, token) +static constexpr int dsv4_hc_post_block_size = 256; + +// comb == nullptr is identity mixing: each destination stream keeps its own residual +// instead of summing across the streams. +template 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, @@ -148,7 +168,7 @@ static void dsv4_hc_post_f32_sycl( 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 block_size = dsv4_hc_post_block_size; const int64_t num_blocks = (nr + block_size - 1) / block_size; stream->parallel_for( @@ -164,8 +184,12 @@ static void dsv4_hc_post_f32_sycl( 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]; + if constexpr (has_comb) { + for (int64_t isrc = 0; isrc < hc; ++isrc) { + sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2]; + } + } else { + sum += residual[i0*sr0 + idst*sr1 + it*sr2]; } dst[i0*sd0 + idst*sd1 + it*sd2] = sum; @@ -189,15 +213,33 @@ void ggml_sycl_op_dsv4_hc_pre(ggml_backend_sycl_context & ctx, ggml_tensor * dst const int64_t hc = x->ne[1]; const int64_t n_tokens = x->ne[2]; + const float scale = ggml_get_op_params_f32(dst, 0); + const bool gated = ggml_get_op_params_i32(dst, 1) != 0; + 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); + if (gated) { + GGML_ASSERT(weights->ne[0] == n_embd); + GGML_ASSERT(weights->ne[1] == hc); + GGML_ASSERT(weights->ne[2] == n_tokens); + 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), nbw2 / sizeof(float), + nbd0 / sizeof(float), nbd1 / sizeof(float), + scale, stream); + } else { + GGML_ASSERT(weights->ne[0] == hc); + GGML_ASSERT(weights->ne[1] == n_tokens); + 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), /*sw2=*/ 0, + nbd0 / sizeof(float), nbd1 / sizeof(float), + scale, stream); + } } void ggml_sycl_op_dsv4_hc_comb(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { @@ -252,24 +294,33 @@ void ggml_sycl_op_dsv4_hc_post(ggml_backend_sycl_context & ctx, ggml_tensor * ds 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); + size_t nbc0 = 0; + size_t nbc1 = 0; + size_t nbc2 = 0; + if (comb) { + GGML_ASSERT(comb->type == GGML_TYPE_F32); + nbc0 = comb->nb[0]; + nbc1 = comb->nb[1]; + nbc2 = comb->nb[2]; + } + 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 auto launch = [&](auto has_comb) { + dsv4_hc_post_f32_sycl( (const float *) x->data, (const float *) residual->data, - (const float *) post->data, (const float *) comb->data, (float *) dst->data, + (const float *) post->data, comb ? (const float *) comb->data : nullptr, (float *) dst->data, n_embd, hc, n_tokens, nbx0 / sizeof(float), nbx1 / sizeof(float), nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float), @@ -277,4 +328,11 @@ void ggml_sycl_op_dsv4_hc_post(ggml_backend_sycl_context & ctx, ggml_tensor * ds nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float), nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float), stream); + }; + + if (comb) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index a7fbd1644c..e599d2d846 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6482,13 +6482,14 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons 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 && ggml_get_op_params_i32(op, 1) == 0; + 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] != nullptr && op->src[3]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && + (op->src[3] == nullptr || 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 &&