char-fork/llama.cpp

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

Akarshan Biswassycl : add flash-attn support for head size 512 (#21654)b54cb2e3d

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.0 KiB59 linesraw
1#include <sycl/sycl.hpp>
2#include <sycl/ext/oneapi/work_group_static.hpp>
3#include "dpct/helper.hpp"
4#include "common.hpp"
5#include "fattn-common.hpp"
6#include "fattn-tile.hpp"
7#include <cmath>
8#include <float.h>
9namespace syclex = sycl::ext::oneapi::experimental;
10
11void ggml_sycl_flash_attn_ext_tile(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
12    const ggml_tensor * K = dst->src[1];
13    const ggml_tensor * V = dst->src[2];
14    switch (K->ne[0]) {
15        case  40: {
16            GGML_ASSERT(V->ne[0] == K->ne[0]);
17            ggml_sycl_flash_attn_ext_tile_case< 40,  40>(ctx, dst);
18        } break;
19        case  64: {
20            GGML_ASSERT(V->ne[0] == K->ne[0]);
21            ggml_sycl_flash_attn_ext_tile_case< 64,  64>(ctx, dst);
22        } break;
23        case  72: {
24            GGML_ASSERT(V->ne[0] == K->ne[0]);
25            ggml_sycl_flash_attn_ext_tile_case< 72,  72>(ctx, dst);
26        } break;
27        case  80: {
28            GGML_ASSERT(V->ne[0] == K->ne[0]);
29            ggml_sycl_flash_attn_ext_tile_case< 80,  80>(ctx, dst);
30        } break;
31        case  96: {
32            GGML_ASSERT(V->ne[0] == K->ne[0]);
33            ggml_sycl_flash_attn_ext_tile_case< 96,  96>(ctx, dst);
34        } break;
35        case 112: {
36            GGML_ASSERT(V->ne[0] == K->ne[0]);
37            ggml_sycl_flash_attn_ext_tile_case<112, 112>(ctx, dst);
38        } break;
39        case 128: {
40            GGML_ASSERT(V->ne[0] == K->ne[0]);
41            ggml_sycl_flash_attn_ext_tile_case<128, 128>(ctx, dst);
42        } break;
43        case 256: {
44            GGML_ASSERT(V->ne[0] == K->ne[0]);
45            ggml_sycl_flash_attn_ext_tile_case<256, 256>(ctx, dst);
46        } break;
47        case 512: {
48            GGML_ASSERT(V->ne[0] == K->ne[0]);
49            ggml_sycl_flash_attn_ext_tile_case<512, 512>(ctx, dst);
50        } break;
51        case 576: {
52            GGML_ASSERT(V->ne[0] == 512);
53            ggml_sycl_flash_attn_ext_tile_case<576, 512>(ctx, dst);
54        } break;
55        default: {
56            GGML_ABORT("Unsupported head size");
57        } break;
58    }
59}