char-fork/llama.cpp

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

Johannes GäßlerHIP: remove rocWMMA FlashAttention (#26046)fa72aeccb

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