char-fork/llama.cpp

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

Pascalcuda: fuse snake activation (mul, sin, sqr, mul, add) (#22667)58e68df0f

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.8 KiB72 linesraw
1#include "snake.cuh"
2#include "convert.cuh"
3
4// Fused Snake activation: y = x + sin^2(a * x) * inv_b
5// x: [T, C] (T contiguous), a: [1, C], inv_b: [1, C]
6// Supports F32, F16, BF16 data with F32 compute.
7
8template <typename T>
9static __global__ void snake_kernel(
10        const T     * __restrict__ x,
11        const float * __restrict__ a,
12        const float * __restrict__ inv_b,
13        T           * __restrict__ dst,
14        const int    total,
15        const uint3  T_len_fastdiv) {
16    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
17    if (idx >= total) return;
18
19    const int c = (int) fastdiv((uint32_t) idx, T_len_fastdiv);
20
21    const float xi = ggml_cuda_cast<float>(x[idx]);
22    const float s  = sinf(a[c] * xi);
23    dst[idx] = ggml_cuda_cast<T>(xi + s * s * inv_b[c]);
24}
25
26// Internal launcher with explicit x/a/inv_b/dst tensors.
27// Shared by the public op (reads dst->src) and the fusion path (explicit args).
28static void launch_snake(ggml_backend_cuda_context & ctx,
29                         const ggml_tensor * x,
30                         const ggml_tensor * a,
31                         const ggml_tensor * inv_b,
32                         ggml_tensor *       dst) {
33    const float * a_d     = (const float *)a->data;
34    const float * inv_b_d = (const float *)inv_b->data;
35
36    const int   T = (int)x->ne[0];
37    const int   C = (int)x->ne[1];
38    const int   total = T * C;
39    const uint3 T_len_fastdiv = init_fastdiv_values((uint64_t) T);
40
41    const int block_size = 256;
42    const int grid_size  = (total + block_size - 1) / block_size;
43
44    cudaStream_t stream = ctx.stream();
45
46    switch (x->type) {
47        case GGML_TYPE_F32: {
48            snake_kernel<<<grid_size, block_size, 0, stream>>>(
49                (const float *)x->data, a_d, inv_b_d, (float *)dst->data, total, T_len_fastdiv);
50        } break;
51        case GGML_TYPE_F16: {
52            snake_kernel<<<grid_size, block_size, 0, stream>>>(
53                (const half *)x->data, a_d, inv_b_d, (half *)dst->data, total, T_len_fastdiv);
54        } break;
55        case GGML_TYPE_BF16: {
56            snake_kernel<<<grid_size, block_size, 0, stream>>>(
57                (const nv_bfloat16 *)x->data, a_d, inv_b_d, (nv_bfloat16 *)dst->data, total, T_len_fastdiv);
58        } break;
59        default:
60            GGML_ABORT("snake: unsupported type");
61    }
62}
63
64// Fusion entry: caller supplies x/a/inv_b explicitly from the matched
65// mul -> sin -> sqr -> mul -> add pattern. The dst is the trailing add output.
66void ggml_cuda_op_snake_fused(ggml_backend_cuda_context & ctx,
67                              const ggml_tensor * x,
68                              const ggml_tensor * a,
69                              const ggml_tensor * inv_b,
70                              ggml_tensor *       dst) {
71    launch_snake(ctx, x, a, inv_b, dst);
72}