char-fork/llama.cpp

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

Johannes GäßlerCUDA: missing PDL sync for FWHT, better fallback (#23690)192d8ae8b

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.1 KiB101 linesraw
1#include "common.cuh"
2#include "fwht.cuh"
3
4template <int N>
5__launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1)
6__global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows, const float scale) {
7    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
8
9    const int64_t r = (int64_t) blockIdx.x * blockDim.y + threadIdx.y;
10
11    if (r >= n_rows) {
12        return;
13    }
14
15    src += r * N;
16    dst += r * N;
17
18    static constexpr int el_w = N / warp_size;
19    float     reg[el_w];
20    const int lane = threadIdx.x;
21
22    ggml_cuda_pdl_sync();
23#pragma unroll
24    for (int i = 0; i < el_w; ++i) {
25        reg[i] = src[i * warp_size + lane] * scale;
26    }
27
28#pragma unroll
29    for (int h = 1; h < warp_size; h *= 2) {
30#pragma unroll
31        for (int j = 0; j < el_w; j++) {
32            const float val  = reg[j];
33            const float val2 = __shfl_xor_sync(0xFFFFFFFF, val, h, warp_size);
34
35            reg[j] = (lane & h) == 0 ? val + val2 : val2 - val;
36        }
37    }
38
39#pragma unroll
40    for (int h = warp_size; h < N; h *= 2) {
41        const int step = h / warp_size;
42#pragma unroll
43        for (int j = 0; j < el_w; j += 2 * step) {
44#pragma unroll
45            for (int k = 0; k < step; k++) {
46                const float x = reg[j + k];
47                const float y = reg[j + k + step];
48
49                reg[j + k]        = x + y;
50                reg[j + k + step] = x - y;
51            }
52        }
53    }
54
55#pragma unroll
56    for (int i = 0; i < el_w; ++i) {
57        dst[i * warp_size + lane] = reg[i];
58    }
59}
60
61bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
62    GGML_ASSERT(ggml_are_same_shape(src, dst));
63    if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
64        return false;
65    }
66    const int     n    = src->ne[0];
67    const int64_t rows = ggml_nrows(src);
68
69    const float * src_d = (const float *) src->data;
70    float *       dst_d = (float *) dst->data;
71
72    const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
73    const int rows_per_block = 4;
74
75    const int64_t num_blocks = (rows + rows_per_block - 1) / rows_per_block;
76
77    cudaStream_t                         stream = ctx.stream();
78    dim3                                 grid_dims(num_blocks, 1, 1);
79    dim3                                 block_dims(warp_size, rows_per_block, 1);
80    const ggml_cuda_kernel_launch_params launch_params =
81        ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream);
82
83    const float scale = 1 / sqrtf(n);
84
85    switch (n) {
86        case 64:
87            ggml_cuda_kernel_launch(fwht_cuda<64>, launch_params, src_d, dst_d, rows, scale);
88            return true;
89        case 128:
90            ggml_cuda_kernel_launch(fwht_cuda<128>, launch_params, src_d, dst_d, rows, scale);
91            return true;
92        case 256:
93            ggml_cuda_kernel_launch(fwht_cuda<256>, launch_params, src_d, dst_d, rows, scale);
94            return true;
95        case 512:
96            ggml_cuda_kernel_launch(fwht_cuda<512>, launch_params, src_d, dst_d, rows, scale);
97            return true;
98        default:
99            return false;
100    }
101}