char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Daniel Beveniusggml : fix padding in timestep embedding kernels (#15932)3913f8730

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.6 KiB73 linesraw
1//
2// MIT license
3// Copyright (C) 2024 Intel Corporation
4// SPDX-License-Identifier: MIT
5//
6
7//
8// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
9// See https://llvm.org/LICENSE.txt for license information.
10// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
11//
12
13#include "tsembd.hpp"
14
15static void timestep_embedding_f32(
16        const float * timesteps, float * dst, const int nb1,
17        const int dim, const int max_period, const sycl::nd_item<3> &item_ct1) {
18    // item_ct1.get_group(1)(blockIDx.y): idx of timesteps->ne[0]
19    // item_ct1.get_group(2) (blockIDx.x): idx of ((dim + 1) / 2) / BLOCK_SIZE
20    int i = item_ct1.get_group(1);
21    int j = item_ct1.get_local_id(2) + item_ct1.get_group(2) * item_ct1.get_local_range(2);
22    float * embed_data = (float *)((char *)dst +  i*nb1);
23
24    int half = dim / 2;
25
26    if (dim % 2 != 0 && j == half) {
27        embed_data[2 * half] = 0.f;
28    }
29
30    if (j >= half) {
31        return;
32    }
33
34    float timestep = timesteps[i];
35    float freq = (float)sycl::native::exp(-(sycl::log((float)max_period)) * j / half);
36    float arg = timestep * freq;
37    embed_data[j] = sycl::cos(arg);
38    embed_data[j + half] = sycl::sin(arg);
39}
40
41static void timestep_embedding_f32_sycl(
42        const float * x, float * dst, const int ne00, const int nb1,
43        const int dim, const int max_period, const queue_ptr& stream) {
44    // As the kernel returns when thread.idx is larger than dim/2, the half_ceil does not need to pad
45    int half_ceil = dim / 2;
46    int num_blocks = (half_ceil + SYCL_TIMESTEP_EMBEDDING_BLOCK_SIZE - 1) / SYCL_TIMESTEP_EMBEDDING_BLOCK_SIZE;
47    sycl::range<3> block_dims(1, 1, SYCL_TIMESTEP_EMBEDDING_BLOCK_SIZE);
48    sycl::range<3> gridDim(1, ne00, num_blocks);
49    stream->parallel_for(
50        sycl::nd_range<3>(
51            gridDim * block_dims, block_dims),
52        [=](sycl::nd_item<3> item_ct1) {
53            timestep_embedding_f32(
54                x, dst, nb1, dim, max_period, item_ct1
55            );
56        });
57}
58
59void ggml_sycl_op_timestep_embedding(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
60    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
61    const ggml_tensor *  src0   = dst->src[0];
62    const float * src0_d = (const float *)src0->data;
63    float * dst_d = (float *)dst->data;
64    dpct::queue_ptr stream = ctx.stream();
65
66    GGML_ASSERT(src0->type == GGML_TYPE_F32);
67    GGML_ASSERT(dst->type == GGML_TYPE_F32);
68
69    const int dim = dst->op_params[0];
70    const int max_period = dst->op_params[1];
71
72    timestep_embedding_f32_sycl(src0_d, dst_d, src0->ne[0], dst->nb[1], dim, max_period, stream);
73}