char-fork/llama.cpp

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

Jingxin (Philip) Lisycl : fix test-backend-ops CI break && restore Kronecker product FWHT support (#28016) (#28254)4d9176092

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.9 KiB291 linesraw
1#include "fwht.hpp"
2
3#include <cmath>
4#define P 1.0f
5#define N -1.0f
6
7// constant Hadamard matrix via Paley I construction
8static constexpr float H12[12][12] = {
9    { P, P, P, P, P, P, P, P, P, P, P, P },
10    { P, N, P, N, P, P, P, N, N, N, P, N },
11    { P, N, N, P, N, P, P, P, N, N, N, P },
12    { P, P, N, N, P, N, P, P, P, N, N, N },
13    { P, N, P, N, N, P, N, P, P, P, N, N },
14    { P, N, N, P, N, N, P, N, P, P, P, N },
15    { P, N, N, N, P, N, N, P, N, P, P, P },
16    { P, P, N, N, N, P, N, N, P, N, P, P },
17    { P, P, P, N, N, N, P, N, N, P, N, P },
18    { P, P, P, P, N, N, N, P, N, N, P, N },
19    { P, N, P, P, P, N, N, N, P, N, N, P },
20    { P, P, N, P, P, P, N, N, N, P, N, N }
21};
22
23static constexpr float H20[20][20] = {
24    { P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P, P },
25    { P, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N },
26    { P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P },
27    { P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P },
28    { P, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N },
29    { P, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N, N },
30    { P, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N, N },
31    { P, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P, N },
32    { P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N, P },
33    { P, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P, N },
34    { P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N, P },
35    { P, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P, N },
36    { P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P, P },
37    { P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P, P },
38    { P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P, P },
39    { P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N, P },
40    { P, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N, N },
41    { P, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P, N },
42    { P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N, P },
43    { P, P, N, N, P, P, P, P, N, P, N, P, N, N, N, N, P, P, N, N }
44};
45
46#undef P
47#undef N
48
49template <int N>
50static void fwht_kernel(const float * __restrict__ src, float * __restrict__ dst, const int64_t n_rows,
51                        const float scale, const sycl::nd_item<2> & item) {
52    const sycl::sub_group sg = item.get_sub_group();
53
54    const int64_t r = item.get_global_id(0);
55    if (r >= n_rows) {
56        return;
57    }
58
59    src += r * N;
60    dst += r * N;
61
62    constexpr int el_w = N / WARP_SIZE;
63    static_assert(el_w >= 1 && N % WARP_SIZE == 0, "row must be a whole number of sub-group widths");
64
65    float     reg[el_w];
66    const int lane = sg.get_local_linear_id();
67
68#pragma unroll
69    for (int i = 0; i < el_w; ++i) {
70        reg[i] = src[i * WARP_SIZE + lane] * scale;
71    }
72
73    // Butterflies inside the sub-group. The partner of a lane with bit h clear is the
74    // lower index of the pair, so it takes the sum and the upper takes lower - upper.
75#pragma unroll
76    for (int h = 1; h < WARP_SIZE; h *= 2) {
77#pragma unroll
78        for (int j = 0; j < el_w; ++j) {
79            const float val  = reg[j];
80            const float val2 = dpct::permute_sub_group_by_xor(sg, val, h, WARP_SIZE);
81
82            reg[j] = (lane & h) == 0 ? val + val2 : val2 - val;
83        }
84    }
85
86    // Butterflies across registers: h is a multiple of WARP_SIZE, so the partner of
87    // element i*WARP_SIZE + lane lives in reg[i + h/WARP_SIZE] on the same lane.
88#pragma unroll
89    for (int h = WARP_SIZE; h < N; h *= 2) {
90        const int step = h / WARP_SIZE;
91#pragma unroll
92        for (int j = 0; j < el_w; j += 2 * step) {
93#pragma unroll
94            for (int k = 0; k < step; ++k) {
95                const float x = reg[j + k];
96                const float y = reg[j + k + step];
97
98                reg[j + k]        = x + y;
99                reg[j + k + step] = x - y;
100            }
101        }
102    }
103
104#pragma unroll
105    for (int i = 0; i < el_w; ++i) {
106        dst[i * WARP_SIZE + lane] = reg[i];
107    }
108}
109
110template <int N>
111static void launch_fwht(const float * src, float * dst, const int64_t n_rows, const float scale,
112                        dpct::queue_ptr stream) {
113    constexpr int rows_per_block = 4;
114
115    const int64_t num_blocks = (n_rows + rows_per_block - 1) / rows_per_block;
116
117    // dim 1 is the fastest-varying, so a sub-group is exactly one row's WARP_SIZE lanes.
118    const sycl::range<2> global(num_blocks * rows_per_block, WARP_SIZE);
119    const sycl::range<2> local(rows_per_block, WARP_SIZE);
120
121    stream->parallel_for(sycl::nd_range<2>(global, local),
122                         [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
123                             fwht_kernel<N>(src, dst, n_rows, scale, item);
124                         });
125}
126
127template <int N, int m>
128static void kronecker_kernel(const float * __restrict__ src,
129                             float * __restrict__ dst,
130                             const int64_t            n_rows,
131                             const float              scale,
132                             const sycl::nd_item<2> & item) {
133    static_assert(m == 12 || m == 20, "block size has to be 12 or 20.");
134
135    const sycl::sub_group sg = item.get_sub_group();
136
137    const int64_t r = item.get_global_id(0);
138    if (r >= n_rows) {
139        return;
140    }
141
142    src += r * N;
143    dst += r * N;
144
145    constexpr int blocks_per_group = N / m;
146    constexpr int el_w             = blocks_per_group / WARP_SIZE;
147    static_assert(el_w >= 1 && blocks_per_group % WARP_SIZE == 0, "blocks_per_group must be a multiple of WARP_SIZE");
148    float     reg[el_w * m];
149    const int lane = sg.get_local_linear_id();
150
151#pragma unroll
152    for (int i = 0; i < el_w; ++i) {
153        const int b_idx = i * WARP_SIZE + lane;
154
155#pragma unroll
156        for (int j = 0; j < m; ++j) {
157            reg[i * m + j] = src[b_idx * m + j] * scale;
158        }
159    }
160
161#pragma unroll
162    for (int b = 0; b < el_w; ++b) {
163        float z[m] = { 0.0f };
164
165#pragma unroll
166        for (int i = 0; i < m; ++i) {
167#pragma unroll
168            for (int j = 0; j < m; ++j) {
169                const float h = (m == 12 ? H12[j][i] : H20[j][i]);
170                z[i] += reg[b * m + j] * h;
171            }
172        }
173
174#pragma unroll
175        for (int i = 0; i < m; ++i) {
176            reg[b * m + i] = z[i];
177        }
178    }
179
180#pragma unroll
181    for (int h = 1; h < WARP_SIZE; h *= 2) {
182#pragma unroll
183        for (int j = 0; j < el_w; ++j) {
184#pragma unroll
185            for (int k = 0; k < m; ++k) {
186                const float val  = reg[j * m + k];
187                const float val2 = dpct::permute_sub_group_by_xor(sg, val, h, WARP_SIZE);
188
189                reg[j * m + k] = (lane & h) == 0 ? val + val2 : val2 - val;
190            }
191        }
192    }
193
194#pragma unroll
195    for (int h = WARP_SIZE; h < blocks_per_group; h *= 2) {
196        const int step = h / WARP_SIZE;
197#pragma unroll
198        for (int j = 0; j < el_w; j += 2 * step) {
199#pragma unroll
200            for (int s = 0; s < step; ++s) {
201#pragma unroll
202                for (int k = 0; k < m; ++k) {
203                    const float x = reg[(j + s) * m + k];
204                    const float y = reg[(j + s + step) * m + k];
205
206                    reg[(j + s) * m + k]        = x + y;
207                    reg[(j + s + step) * m + k] = x - y;
208                }
209            }
210        }
211    }
212
213#pragma unroll
214    for (int i = 0; i < el_w; ++i) {
215        const int b_idx = i * WARP_SIZE + lane;
216#pragma unroll
217        for (int k = 0; k < m; ++k) {
218            dst[b_idx * m + k] = reg[i * m + k];
219        }
220    }
221}
222
223template <int N, int m>
224static void launch_kronecker(const float *   src,
225                             float *         dst,
226                             const int64_t   n_rows,
227                             const float     scale,
228                             dpct::queue_ptr stream) {
229    constexpr int rows_per_block = 4;
230
231    const int64_t num_blocks = (n_rows + rows_per_block - 1) / rows_per_block;
232
233    // dim 1 is the fastest-varying, so a sub-group is exactly one row's WARP_SIZE lanes.
234    const sycl::range<2> global(num_blocks * rows_per_block, WARP_SIZE);
235    const sycl::range<2> local(rows_per_block, WARP_SIZE);
236
237    stream->parallel_for(sycl::nd_range<2>(global, local),
238                         [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
239                             kronecker_kernel<N, m>(src, dst, n_rows, scale, item);
240                         });
241}
242
243bool ggml_sycl_op_fwht(ggml_backend_sycl_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
244    if (src->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
245        return false;
246    }
247    if (!ggml_are_same_shape(src, dst)) {
248        return false;
249    }
250    if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
251        return false;
252    }
253
254    const int     n    = (int) src->ne[0];
255    const int64_t rows = ggml_nrows(src);
256
257    const float *   src_d  = (const float *) src->data;
258    float *         dst_d  = (float *) dst->data;
259    dpct::queue_ptr stream = ctx.stream();
260
261    const float scale = 1.0f / std::sqrt((float) n);
262
263    switch (n) {
264        case 64:
265            launch_fwht<64>(src_d, dst_d, rows, scale, stream);
266            return true;
267        case 128:
268            launch_fwht<128>(src_d, dst_d, rows, scale, stream);
269            return true;
270        case 256:
271            launch_fwht<256>(src_d, dst_d, rows, scale, stream);
272            return true;
273        case 512:
274            launch_fwht<512>(src_d, dst_d, rows, scale, stream);
275            return true;
276        case 384:
277            launch_kronecker<384, 12>(src_d, dst_d, rows, scale, stream);
278            return true;
279        case 768:
280            launch_kronecker<768, 12>(src_d, dst_d, rows, scale, stream);
281            return true;
282        case 640:
283            launch_kronecker<640, 20>(src_d, dst_d, rows, scale, stream);
284            return true;
285        case 1280:
286            launch_kronecker<1280, 20>(src_d, dst_d, rows, scale, stream);
287            return true;
288        default:
289            return false;
290    }
291}