From 6703d7894c70e8b076ce4608157d056e42e6889c Mon Sep 17 00:00:00 2001 From: Frosty40 Date: Thu, 3 Sep 2026 23:05:40 -0500 Subject: [PATCH] sycl: fuse rms_norm+mul+add and add+add residual chains (#27610) Fuse RMS_NORM+MUL+ADD and ADD+ADD under GGML_SYCL_ENABLE_FUSION. ADD+ADD uses the same binbcast indexing and type matrix as standalone add() (f32, f16, f16/f32, i32, i16, bf16, including broadcast and non-contiguous). Unsupported combinations fall back to two add() launches. --- docs/backend/SYCL.md | 2 +- ggml/src/ggml-sycl/binbcast.cpp | 292 +++++++++++++++++++++++++++++++ ggml/src/ggml-sycl/binbcast.hpp | 30 ++++ ggml/src/ggml-sycl/fusion.cpp | 45 ++++- ggml/src/ggml-sycl/ggml-sycl.cpp | 12 ++ ggml/src/ggml-sycl/norm.cpp | 151 +++++++++++++++- ggml/src/ggml-sycl/norm.hpp | 2 + tests/test-backend-ops.cpp | 64 +++++++ 8 files changed, 593 insertions(+), 5 deletions(-) diff --git a/docs/backend/SYCL.md b/docs/backend/SYCL.md index b30e109ad..c4dcb02ee 100644 --- a/docs/backend/SYCL.md +++ b/docs/backend/SYCL.md @@ -806,7 +806,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm | GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` | | GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. | | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. | -| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. | +| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. Unsupported types and layouts fall back to the standalone op kernels. See `ggml_sycl_can_fuse()`. | | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. | | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.
Recommended to use when --split-mode = layer | | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. | diff --git a/ggml/src/ggml-sycl/binbcast.cpp b/ggml/src/ggml-sycl/binbcast.cpp index 306eeddc0..f2f7c4cde 100644 --- a/ggml/src/ggml-sycl/binbcast.cpp +++ b/ggml/src/ggml-sycl/binbcast.cpp @@ -1,5 +1,6 @@ #include "binbcast.hpp" +#include #include #include #include @@ -356,3 +357,294 @@ void ggml_sycl_repeat(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { ggml_sycl_op_repeat(ctx, dst); } +// fused ADD+ADD: dst = (src0 + src1) + src2. Same indexing as k_bin_bcast, so mixed +// types, broadcast, and non-contiguous layouts that add() already handles also fuse. +template +static void k_bin_bcast3(const src0_t * src0, const src1_t * src1, const src2_t * src2, dst_t * dst, + int ne0, int ne1, int ne2, int ne3, + int ne10, int ne11, int ne12, int ne13, + int ne20, int ne21, int ne22, int ne23, + int s1, int s2, int s3, + int s00, int s01, int s02, int s03, + int s10, int s11, int s12, int s13, + int s20, int s21, int s22, int s23, + const sycl::nd_item<3> & item_ct1) { + const int i0s = item_ct1.get_local_range(2) * item_ct1.get_group(2) + + item_ct1.get_local_id(2); + const int i1 = (item_ct1.get_local_range(1) * item_ct1.get_group(1) + + item_ct1.get_local_id(1)); + const int i2 = (item_ct1.get_local_range(0) * item_ct1.get_group(0) + + item_ct1.get_local_id(0)) / + ne3; + const int i3 = (item_ct1.get_local_range(0) * item_ct1.get_group(0) + + item_ct1.get_local_id(0)) % + ne3; + + if (i0s >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) { + return; + } + + const int i11 = i1 % ne11; + const int i12 = i2 % ne12; + const int i13 = i3 % ne13; + const int i21 = i1 % ne21; + const int i22 = i2 % ne22; + const int i23 = i3 % ne23; + + const size_t i_src0 = i3 * s03 + i2 * s02 + i1 * s01; + const size_t i_src1 = i13 * s13 + i12 * s12 + i11 * s11; + const size_t i_src2 = i23 * s23 + i22 * s22 + i21 * s21; + const size_t i_dst = i3 * s3 + i2 * s2 + i1 * s1; + + const src0_t * src0_row = src0 + i_src0; + const src1_t * src1_row = src1 + i_src1; + const src2_t * src2_row = src2 + i_src2; + dst_t * dst_row = dst + i_dst; + + for (int i0 = i0s; i0 < ne0; + i0 += item_ct1.get_local_range(2) * item_ct1.get_group_range(2)) { + const int i10 = i0 % ne10; + const int i20 = i0 % ne20; + const float acc = bin_op((float) src0_row[i0 * s00], (float) src1_row[i10 * s10]); + dst_row[i0] = (dst_t) bin_op(acc, (float) src2_row[i20 * s20]); + } +} + +template +static void k_bin_bcast3_unravel(const src0_t * src0, const src1_t * src1, const src2_t * src2, dst_t * dst, + int ne0, int ne1, int ne2, int ne3, + int ne10, int ne11, int ne12, int ne13, + int ne20, int ne21, int ne22, int ne23, + int s1, int s2, int s3, + int s00, int s01, int s02, int s03, + int s10, int s11, int s12, int s13, + int s20, int s21, int s22, int s23, + const sycl::nd_item<3> & item_ct1) { + const int i = item_ct1.get_local_range(2) * item_ct1.get_group(2) + + item_ct1.get_local_id(2); + + const int i3 = i / (ne2 * ne1 * ne0); + const int i2 = (i / (ne1 * ne0)) % ne2; + const int i1 = (i / ne0) % ne1; + const int i0 = i % ne0; + + if (i0 >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) { + return; + } + + const int i11 = i1 % ne11; + const int i12 = i2 % ne12; + const int i13 = i3 % ne13; + const int i21 = i1 % ne21; + const int i22 = i2 % ne22; + const int i23 = i3 % ne23; + + const size_t i_src0 = i3 * s03 + i2 * s02 + i1 * s01; + const size_t i_src1 = i13 * s13 + i12 * s12 + i11 * s11; + const size_t i_src2 = i23 * s23 + i22 * s22 + i21 * s21; + const size_t i_dst = i3 * s3 + i2 * s2 + i1 * s1; + + const int i10 = i0 % ne10; + const int i20 = i0 % ne20; + const float acc = bin_op((float) src0[i_src0 + i0 * s00], (float) src1[i_src1 + i10 * s10]); + dst[i_dst + i0] = (dst_t) bin_op(acc, (float) src2[i_src2 + i20 * s20]); +} + +template +static void launch_bin_bcast3(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + const ggml_tensor * src2, ggml_tensor * dst) { + dpct::queue_ptr stream = ctx.stream(); + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); + + GGML_TENSOR_TERNARY_OP_LOCALS + + int nr1[4] = { (int) (ne10 / ne0), (int) (ne11 / ne1), (int) (ne12 / ne2), (int) (ne13 / ne3) }; + int nr2[4] = { (int) (ne20 / ne0), (int) (ne21 / ne1), (int) (ne22 / ne2), (int) (ne23 / ne3) }; + + int64_t cne[] = { ne0, ne1, ne2, ne3 }; + int64_t cne0[] = { ne00, ne01, ne02, ne03 }; + int64_t cne1[] = { ne10, ne11, ne12, ne13 }; + int64_t cne2[] = { ne20, ne21, ne22, ne23 }; + size_t cnb[] = { nb0, nb1, nb2, nb3 }; + size_t cnb0[] = { nb00, nb01, nb02, nb03 }; + size_t cnb1[] = { nb10, nb11, nb12, nb13 }; + size_t cnb2[] = { nb20, nb21, nb22, nb23 }; + + auto collapse = [](int64_t cne[]) { + cne[0] *= cne[1]; + cne[1] = cne[2]; + cne[2] = cne[3]; + cne[3] = 1; + }; + + auto collapse_nb = [](size_t cnb[], int64_t cne[]) { + cnb[1] *= cne[1]; + cnb[2] *= cne[2]; + cnb[3] *= cne[3]; + }; + + const bool can_collapse = ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(src2) && + !ggml_is_permuted(src0) && !ggml_is_permuted(src1) && !ggml_is_permuted(src2); + if (can_collapse) { + for (int i = 0; i < 4; i++) { + if (nr1[i] != 1 || nr2[i] != 1) { + break; + } + if (i > 0) { + collapse_nb(cnb, cne); + collapse_nb(cnb0, cne0); + collapse_nb(cnb1, cne1); + collapse_nb(cnb2, cne2); + collapse(cne); + collapse(cne0); + collapse(cne1); + collapse(cne2); + } + } + } + + { + int64_t ne0 = cne[0]; + int64_t ne1 = cne[1]; + int64_t ne2 = cne[2]; + int64_t ne3 = cne[3]; + + int64_t ne10 = cne1[0]; + int64_t ne11 = cne1[1]; + int64_t ne12 = cne1[2]; + int64_t ne13 = cne1[3]; + + int64_t ne20 = cne2[0]; + int64_t ne21 = cne2[1]; + int64_t ne22 = cne2[2]; + int64_t ne23 = cne2[3]; + + size_t s1 = cnb[1] / sizeof(dst_t); + size_t s2 = cnb[2] / sizeof(dst_t); + size_t s3 = cnb[3] / sizeof(dst_t); + + size_t s00 = cnb0[0] / sizeof(src0_t); + size_t s01 = cnb0[1] / sizeof(src0_t); + size_t s02 = cnb0[2] / sizeof(src0_t); + size_t s03 = cnb0[3] / sizeof(src0_t); + + size_t s10 = cnb1[0] / sizeof(src1_t); + size_t s11 = cnb1[1] / sizeof(src1_t); + size_t s12 = cnb1[2] / sizeof(src1_t); + size_t s13 = cnb1[3] / sizeof(src1_t); + + size_t s20 = cnb2[0] / sizeof(src2_t); + size_t s21 = cnb2[1] / sizeof(src2_t); + size_t s22 = cnb2[2] / sizeof(src2_t); + size_t s23 = cnb2[3] / sizeof(src2_t); + + GGML_ASSERT(cnb[0] % sizeof(dst_t) == 0 && cnb[1] % sizeof(dst_t) == 0 && cnb[2] % sizeof(dst_t) == 0 && + cnb[3] % sizeof(dst_t) == 0); + GGML_ASSERT(cnb0[0] % sizeof(src0_t) == 0 && cnb0[1] % sizeof(src0_t) == 0 && cnb0[2] % sizeof(src0_t) == 0 && + cnb0[3] % sizeof(src0_t) == 0); + GGML_ASSERT(cnb1[0] % sizeof(src1_t) == 0 && cnb1[1] % sizeof(src1_t) == 0 && cnb1[2] % sizeof(src1_t) == 0 && + cnb1[3] % sizeof(src1_t) == 0); + GGML_ASSERT(cnb2[0] % sizeof(src2_t) == 0 && cnb2[1] % sizeof(src2_t) == 0 && cnb2[2] % sizeof(src2_t) == 0 && + cnb2[3] % sizeof(src2_t) == 0); + + const src0_t * src0_dd = (const src0_t *) src0->data; + const src1_t * src1_dd = (const src1_t *) src1->data; + const src2_t * src2_dd = (const src2_t *) src2->data; + dst_t * dst_dd = (dst_t *) dst->data; + + const int block_size = 128; + int64_t hne0 = std::max(ne0 / 2LL, 1LL); + + sycl::range<3> block_dims(1, 1, 1); + block_dims[2] = std::min(hne0, block_size); + block_dims[1] = std::min(ne1, block_size / (unsigned int) block_dims[2]); + block_dims[0] = std::min(std::min(ne2 * ne3, + block_size / (unsigned int) block_dims[2] / + (unsigned int) block_dims[1]), + 64U); + + sycl::range<3> block_nums((ne2 * ne3 + block_dims[0] - 1) / block_dims[0], + (ne1 + block_dims[1] - 1) / block_dims[1], + (hne0 + block_dims[2] - 1) / block_dims[2]); + + dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 }); + + if (block_nums[0] > 65535) { + int block_num = (ne0 * ne1 * ne2 * ne3 + block_size - 1) / block_size; + stream->parallel_for( + sycl::nd_range<3>(sycl::range<3>(1, 1, block_num) * sycl::range<3>(1, 1, block_size), + sycl::range<3>(1, 1, block_size)), + [=](sycl::nd_item<3> item_ct1) { + k_bin_bcast3_unravel(src0_dd, src1_dd, src2_dd, dst_dd, ne0, ne1, ne2, ne3, ne10, ne11, + ne12, ne13, ne20, ne21, ne22, ne23, s1, s2, s3, s00, s01, s02, s03, + s10, s11, s12, s13, s20, s21, s22, s23, item_ct1); + }); + } else { + stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims), + [=](sycl::nd_item<3> item_ct1) { + k_bin_bcast3(src0_dd, src1_dd, src2_dd, dst_dd, ne0, ne1, ne2, ne3, ne10, + ne11, ne12, ne13, ne20, ne21, ne22, ne23, s1, s2, s3, s00, + s01, s02, s03, s10, s11, s12, s13, s20, s21, s22, s23, + item_ct1); + }); + } + } +} + +void ggml_sycl_op_add_add_fused(ggml_backend_sycl_context & ctx, ggml_tensor * add0, ggml_tensor * add1) { + const ggml_tensor * src0 = add0->src[0]; + const ggml_tensor * src1 = add0->src[1]; + const ggml_tensor * src2 = add1->src[1]; + ggml_tensor * dst = add1; + + GGML_ASSERT(add1->src[0] == add0); + GGML_ASSERT(ggml_sycl_add_kernel_supports(src0->type, src1->type, add0->type)); + GGML_ASSERT(ggml_sycl_add_kernel_supports(add0->type, src2->type, dst->type)); + + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 && + dst->type == GGML_TYPE_F32) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && src2->type == GGML_TYPE_F16 && + dst->type == GGML_TYPE_F16) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 && + dst->type == GGML_TYPE_F16) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && src2->type == GGML_TYPE_F32 && + dst->type == GGML_TYPE_F16) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F16 && + dst->type == GGML_TYPE_F16) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32 && src2->type == GGML_TYPE_I32 && + dst->type == GGML_TYPE_I32) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_I16 && src1->type == GGML_TYPE_I16 && src2->type == GGML_TYPE_I16 && + dst->type == GGML_TYPE_I16) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); +#ifdef GGML_SYCL_HAS_BF16 + } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && src2->type == GGML_TYPE_BF16 && + dst->type == GGML_TYPE_BF16) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 && + dst->type == GGML_TYPE_BF16) { + launch_bin_bcast3( + ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && src2->type == GGML_TYPE_F32 && + dst->type == GGML_TYPE_BF16) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); + } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_BF16 && + dst->type == GGML_TYPE_BF16) { + launch_bin_bcast3(ctx, src0, src1, src2, dst); +#endif + } else { + fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s, src1: %s, src2: %s\n", __func__, + ggml_type_name(dst->type), ggml_type_name(src0->type), ggml_type_name(src1->type), + ggml_type_name(src2->type)); + GGML_ABORT("fatal error"); + } +} + diff --git a/ggml/src/ggml-sycl/binbcast.hpp b/ggml/src/ggml-sycl/binbcast.hpp index 9cce0f053..0e5a5ca1c 100644 --- a/ggml/src/ggml-sycl/binbcast.hpp +++ b/ggml/src/ggml-sycl/binbcast.hpp @@ -34,6 +34,36 @@ void ggml_sycl_div(ggml_backend_sycl_context & ctx, ggml_tensor * dst); void ggml_sycl_repeat(ggml_backend_sycl_context & ctx, ggml_tensor * dst); +void ggml_sycl_op_add_add_fused(ggml_backend_sycl_context & ctx, ggml_tensor * add0, ggml_tensor * add1); + +// Type combinations the standalone SYCL add() kernel can run. Fused ADD+ADD +// uses the same set; anything else falls back to two add() launches. +inline bool ggml_sycl_add_kernel_supports(enum ggml_type src0, enum ggml_type src1, enum ggml_type dst) { + if (src0 == GGML_TYPE_F32 && src1 == GGML_TYPE_F32 && dst == GGML_TYPE_F32) { + return true; + } + if (src0 == GGML_TYPE_F16 && src1 == GGML_TYPE_F16 && dst == GGML_TYPE_F16) { + return true; + } + if (src0 == GGML_TYPE_F16 && src1 == GGML_TYPE_F32 && dst == GGML_TYPE_F16) { + return true; + } + if (src0 == GGML_TYPE_I32 && src1 == GGML_TYPE_I32 && dst == GGML_TYPE_I32) { + return true; + } + if (src0 == GGML_TYPE_I16 && src1 == GGML_TYPE_I16 && dst == GGML_TYPE_I16) { + return true; + } +#ifdef GGML_SYCL_HAS_BF16 + if (src0 == GGML_TYPE_BF16 && src1 == GGML_TYPE_BF16 && dst == GGML_TYPE_BF16) { + return true; + } + if (src0 == GGML_TYPE_BF16 && src1 == GGML_TYPE_F32 && dst == GGML_TYPE_BF16) { + return true; + } +#endif + return false; +} #endif //GGML_SYCL_BINBCAST_HPP diff --git a/ggml/src/ggml-sycl/fusion.cpp b/ggml/src/ggml-sycl/fusion.cpp index 709bc8ca2..b5e79bea5 100644 --- a/ggml/src/ggml-sycl/fusion.cpp +++ b/ggml/src/ggml-sycl/fusion.cpp @@ -1,4 +1,5 @@ #include "fusion.hpp" +#include "binbcast.hpp" #include @@ -94,9 +95,14 @@ bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializ return false; } - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { + if ((ops.size() == 2 || ops.size() == 3) && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { + if (ops.size() == 3 && ops.begin()[2] != GGML_OP_ADD) { + return false; + } + const ggml_tensor * rms_norm = cgraph->nodes[node_idx]; const ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + const ggml_tensor * add = ops.size() == 3 ? cgraph->nodes[node_idx + 2] : nullptr; GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32); GGML_ASSERT(rms_norm->type == GGML_TYPE_F32); @@ -122,6 +128,43 @@ bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializ return false; } + if (add != nullptr) { + if (add->src[0]->type != GGML_TYPE_F32 || + add->src[1]->type != GGML_TYPE_F32 || + add->type != GGML_TYPE_F32) { + return false; + } + + // the fused kernel indexes the residual as add[col] and does not broadcast it + const ggml_tensor * add_w = (add->src[0] == mul) ? add->src[1] : add->src[0]; + if (!ggml_are_same_shape(add_w, add)) { + return false; + } + + if (!ggml_is_contiguous(add->src[0]) || !ggml_is_contiguous_rows(add->src[1])) { + return false; + } + } + + return true; + } + + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_ADD && ops.begin()[1] == GGML_OP_ADD) { + const ggml_tensor * add0 = cgraph->nodes[node_idx]; + const ggml_tensor * add1 = cgraph->nodes[node_idx + 1]; + // ggml_can_fuse already guarantees add1 consumes add0 and that add0 has a single use. + // Keep the CUDA association: the running sum is src0 of the next ADD so the fused + // float fold matches two sequential add() launches. + if (add1->src[0] != add0) { + return false; + } + + const ggml_tensor * c = add1->src[1]; + if (!ggml_sycl_add_kernel_supports(add0->src[0]->type, add0->src[1]->type, add0->type) || + !ggml_sycl_add_kernel_supports(add0->type, c->type, add1->type)) { + return false; + } + return true; } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index feced748c..27804e073 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -5865,12 +5865,24 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc continue; } } + if (node->op == GGML_OP_RMS_NORM && + ggml_sycl_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { + ggml_sycl_op_rms_norm_fused_add(*sycl_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); + i += 2; + continue; + } if (node->op == GGML_OP_RMS_NORM && ggml_sycl_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { ggml_sycl_op_rms_norm_fused(*sycl_ctx, node, cgraph->nodes[i + 1]); i++; continue; } + if (node->op == GGML_OP_ADD && + ggml_sycl_can_fuse(cgraph, i, { GGML_OP_ADD, GGML_OP_ADD }, {})) { + ggml_sycl_op_add_add_fused(*sycl_ctx, node, cgraph->nodes[i + 1]); + i++; + continue; + } if (node->op == GGML_OP_UNARY && ggml_sycl_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { ggml_get_unary_op(node) })) { ggml_sycl_op_unary_mul_fused(*sycl_ctx, node, cgraph->nodes[i + 1]); diff --git a/ggml/src/ggml-sycl/norm.cpp b/ggml/src/ggml-sycl/norm.cpp index f98a7a954..2d3033729 100644 --- a/ggml/src/ggml-sycl/norm.cpp +++ b/ggml/src/ggml-sycl/norm.cpp @@ -144,13 +144,17 @@ static void group_norm_f32(const float* x, float* dst, const int group_size, con } } -template +template static void rms_norm_f32(const float* x, float* dst, const int ncols, const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample, const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, int block_size, const float* mul = nullptr, const int64_t mul_stride_row = 0, const int64_t mul_stride_channel = 0, - const int64_t mul_stride_sample = 0, const int mul_nrows = 0, const int mul_nchannels = 0, const int mul_nsamples = 0) { + const int64_t mul_stride_sample = 0, const int mul_nrows = 0, const int mul_nchannels = 0, const int mul_nsamples = 0, + const float* add = nullptr, const int64_t add_stride_row = 0, const int64_t add_stride_channel = 0, + const int64_t add_stride_sample = 0, const int add_nrows = 0, const int add_nchannels = 0, const int add_nsamples = 0) { + + static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying"); const int sample = item_ct1.get_group(0); const int channel = item_ct1.get_group(1); @@ -174,6 +178,13 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols, mul += mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row; } + if constexpr (do_add) { + const int add_row = row % add_nrows; + const int add_channel = channel % add_nchannels; + const int add_sample = sample % add_nsamples; + add += add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row; + } + float tmp = 0.0f; // partial sum for thread in warp for (int col = tid; col < ncols; col += block_size) { @@ -205,7 +216,9 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols, const float scale = sycl::rsqrt(mean + eps); for (int col = tid; col < ncols; col += block_size) { - if constexpr (do_multiply) { + if constexpr (do_multiply && do_add) { + dst[col * dst_stride_col] = scale * x[col * src_stride_col] * mul[col] + add[col]; + } else if constexpr (do_multiply) { dst[col * dst_stride_col] = scale * x[col * src_stride_col] * mul[col]; } else { dst[col * dst_stride_col] = scale * x[col * src_stride_col]; @@ -424,6 +437,53 @@ static void rms_norm_mul_f32_sycl(const float* x, const float* mul, float* dst, } } +static void rms_norm_mul_add_f32_sycl(const float* x, const float* mul, const float* add, float* dst, + const int ncols, const int nrows, const int nchannels, const int nsamples, + const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample, + const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, + const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample, + const int mul_nrows, const int mul_nchannels, const int mul_nsamples, + const int64_t add_stride_row, const int64_t add_stride_channel, const int64_t add_stride_sample, + const int add_nrows, const int add_nchannels, const int add_nsamples, + const float eps, queue_ptr stream, int device) { + const sycl::range<3> global_dims(nsamples, nchannels, nrows); + if (ncols < 1024) { + const sycl::range<3> block_dims(1, 1, WARP_SIZE); + stream->submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<3>(global_dims * block_dims, block_dims), + [=](sycl::nd_item<3> item_ct1) + [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + rms_norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, nullptr, WARP_SIZE, + mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_nrows, mul_nchannels, mul_nsamples, + add, add_stride_row, add_stride_channel, add_stride_sample, add_nrows, add_nchannels, add_nsamples); + }); + }); + } + else { + const int work_group_size = ggml_sycl_info().max_work_group_sizes[device]; + assert(work_group_size % (WARP_SIZE * WARP_SIZE) == 0); + const sycl::range<3> block_dims(1, 1, work_group_size); + stream->submit([&](sycl::handler& cgh) { + sycl::local_accessor s_sum_acc_ct1(sycl::range<1>(work_group_size / WARP_SIZE), cgh); + cgh.parallel_for( + sycl::nd_range<3>(global_dims * block_dims, block_dims), + [=](sycl::nd_item<3> item_ct1) + [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + rms_norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, get_pointer(s_sum_acc_ct1), work_group_size, + mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_nrows, mul_nchannels, mul_nsamples, + add, add_stride_row, add_stride_channel, add_stride_sample, add_nrows, add_nchannels, add_nsamples); + }); + }); + } +} + template static void l2_norm_f32_sycl(const float * x, float * dst, @@ -626,6 +686,91 @@ void ggml_sycl_op_rms_norm_fused(ggml_backend_sycl_context & ctx, ggml_tensor * mul_s01, mul_s02, mul_s03, mul_nrows, mul_nchannels, mul_nsamples, eps, main_stream, ctx.device); } +void ggml_sycl_op_rms_norm_fused_add(ggml_backend_sycl_context & ctx, ggml_tensor * dst, + ggml_tensor * mul_tensor, ggml_tensor * add_tensor) { + const ggml_tensor * rms_norm_src = dst->src[0]; + float eps = 0.0f; + memcpy(&eps, dst->op_params, sizeof(float)); + + const float * src0_dd = static_cast(rms_norm_src->data); + const float * mul_dd = nullptr; + const ggml_tensor * mul_src = nullptr; + if (mul_tensor->src[0] == dst) { + mul_dd = static_cast(mul_tensor->src[1]->data); + mul_src = mul_tensor->src[1]; + } else if (mul_tensor->src[1] == dst) { + mul_dd = static_cast(mul_tensor->src[0]->data); + mul_src = mul_tensor->src[0]; + } else { + GGML_ASSERT(false); + } + + const float * add_dd = nullptr; + const ggml_tensor * add_src = nullptr; + if (add_tensor->src[0] == mul_tensor) { + add_dd = static_cast(add_tensor->src[1]->data); + add_src = add_tensor->src[1]; + } else if (add_tensor->src[1] == mul_tensor) { + add_dd = static_cast(add_tensor->src[0]->data); + add_src = add_tensor->src[0]; + } else { + GGML_ASSERT(false); + } + + float * dst_dd = static_cast(add_tensor->data); + + dpct::queue_ptr main_stream = ctx.stream(); + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); + + GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(add_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(eps >= 0.0f); + + const int64_t ne00 = rms_norm_src->ne[0]; + const int64_t ne01 = rms_norm_src->ne[1]; + const int64_t ne02 = rms_norm_src->ne[2]; + const int64_t ne03 = rms_norm_src->ne[3]; + + const size_t ts0 = ggml_type_size(rms_norm_src->type); + GGML_ASSERT(rms_norm_src->nb[0] == ts0); + const int64_t s00 = rms_norm_src->nb[0] / ts0; + const int64_t s01 = rms_norm_src->nb[1] / ts0; + const int64_t s02 = rms_norm_src->nb[2] / ts0; + const int64_t s03 = rms_norm_src->nb[3] / ts0; + + const size_t tdst = ggml_type_size(add_tensor->type); + GGML_ASSERT(add_tensor->nb[0] == tdst); + const int64_t d00 = add_tensor->nb[0] / tdst; + const int64_t d01 = add_tensor->nb[1] / tdst; + const int64_t d02 = add_tensor->nb[2] / tdst; + const int64_t d03 = add_tensor->nb[3] / tdst; + + const size_t ts_mul = ggml_type_size(mul_src->type); + GGML_ASSERT(mul_src->nb[0] == ts_mul); + const int64_t mul_s01 = mul_src->nb[1] / ts_mul; + const int64_t mul_s02 = mul_src->nb[2] / ts_mul; + const int64_t mul_s03 = mul_src->nb[3] / ts_mul; + const int mul_nrows = mul_src->ne[1]; + const int mul_nchannels = mul_src->ne[2]; + const int mul_nsamples = mul_src->ne[3]; + + const size_t ts_add = ggml_type_size(add_src->type); + GGML_ASSERT(add_src->nb[0] == ts_add); + const int64_t add_s01 = add_src->nb[1] / ts_add; + const int64_t add_s02 = add_src->nb[2] / ts_add; + const int64_t add_s03 = add_src->nb[3] / ts_add; + const int add_nrows = add_src->ne[1]; + const int add_nchannels = add_src->ne[2]; + const int add_nsamples = add_src->ne[3]; + + rms_norm_mul_add_f32_sycl(src0_dd, mul_dd, add_dd, dst_dd, ne00, ne01, ne02, ne03, + s00, s01, s02, s03, d00, d01, d02, d03, + mul_s01, mul_s02, mul_s03, mul_nrows, mul_nchannels, mul_nsamples, + add_s01, add_s02, add_s03, add_nrows, add_nchannels, add_nsamples, eps, main_stream, ctx.device); +} + void ggml_sycl_op_rms_norm_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2); diff --git a/ggml/src/ggml-sycl/norm.hpp b/ggml/src/ggml-sycl/norm.hpp index 51217c421..ef7b2d386 100644 --- a/ggml/src/ggml-sycl/norm.hpp +++ b/ggml/src/ggml-sycl/norm.hpp @@ -21,6 +21,8 @@ void ggml_sycl_op_rms_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst); void ggml_sycl_op_rms_norm_fused(ggml_backend_sycl_context& ctx, ggml_tensor* dst, ggml_tensor* mul); +void ggml_sycl_op_rms_norm_fused_add(ggml_backend_sycl_context& ctx, ggml_tensor* dst, ggml_tensor* mul_tensor, ggml_tensor* add_tensor); + void ggml_sycl_op_rms_norm_back(ggml_backend_sycl_context& ctx, ggml_tensor* dst); void ggml_sycl_op_group_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index eff63575f..2cf9d9caf 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3700,6 +3700,60 @@ struct test_rms_norm_mul_add : public test_case { } }; +// GGML_OP_ADD + GGML_OP_ADD (fused residual chain) +struct test_add_add : public test_case { + const ggml_type type; + const ggml_type type_addend; + const std::array ne; + const bool broadcast; + const bool view; // non-contiguous a via view_4d + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "ADD_ADD"; + } + + bool run_whole_graph() override { return true; } + + std::string vars() override { + return VARS_TO_STR5(type, type_addend, ne, broadcast, view); + } + + test_add_add(ggml_type type = GGML_TYPE_F32, + ggml_type type_addend = GGML_TYPE_F32, + std::array ne = {64, 5, 4, 3}, + bool broadcast = false, + bool view = false) + : type(type), type_addend(type_addend), ne(ne), broadcast(broadcast), view(view) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + std::array broadcast_dims = {ne[0], 1, 1, 1}; + + ggml_tensor * a; + if (view) { + std::array parent = { ne[0] * 3, ne[1] * 2, ne[2], ne[3] }; + a = ggml_new_tensor(ctx, type, 4, parent.data()); + ggml_set_name(a, "a_parent"); + a = ggml_view_4d(ctx, a, ne[0], ne[1], ne[2], ne[3], a->nb[1], a->nb[2], a->nb[3], 0); + ggml_set_name(a, "a"); + } else { + a = ggml_new_tensor(ctx, type, 4, ne.data()); + ggml_set_name(a, "a"); + } + + ggml_tensor * b = ggml_new_tensor(ctx, type_addend, 4, ne.data()); + ggml_tensor * c = ggml_new_tensor(ctx, type_addend, 4, broadcast ? broadcast_dims.data() : ne.data()); + + ggml_set_name(b, "b"); + ggml_set_name(c, "c"); + + ggml_tensor * out = ggml_add(ctx, ggml_add(ctx, a, b), c); + ggml_set_name(out, "out"); + + return out; + } +}; + // GGML_OP_ADD + GGML_OP_RMS_NORM (fused operation) struct test_add_rms_norm : public test_case { const ggml_type type; @@ -9259,6 +9313,8 @@ static std::vector> make_test_cases_eval() { // fusion test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 5, 4, 3}, {2, 1, 1, 1}, 2)); + test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F16, {10, 5, 4, 3}, {2, 1, 1, 1}, 2)); + test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 1, 1, 1}, 2, true)); test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 2, 1, 1}, 3)); test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 5, 4, 3}, {1, 1, 2, 1}, 4)); test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 1, 1, 2}, 5)); @@ -9314,6 +9370,14 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_add_rms_norm(GGML_TYPE_F32, {n, 1, 1, 1}, 1e-6f, false)); } + for (uint32_t n : {64, 1025}) { + test_cases.emplace_back(new test_add_add(GGML_TYPE_F32, GGML_TYPE_F32, { n, 5, 4, 3 }, false, false)); + test_cases.emplace_back(new test_add_add(GGML_TYPE_F32, GGML_TYPE_F32, { n, 5, 4, 3 }, true, false)); + test_cases.emplace_back(new test_add_add(GGML_TYPE_F32, GGML_TYPE_F32, { n, 5, 4, 3 }, false, true)); + test_cases.emplace_back(new test_add_add(GGML_TYPE_F16, GGML_TYPE_F16, { n, 5, 4, 3 }, false, false)); + test_cases.emplace_back(new test_add_add(GGML_TYPE_F16, GGML_TYPE_F32, { n, 5, 4, 3 }, false, false)); + test_cases.emplace_back(new test_add_add(GGML_TYPE_F16, GGML_TYPE_F32, { n, 5, 4, 3 }, true, false)); + } for (auto multi_add : {false, true}) { for (auto set_rows : {false, true}) {