char-fork/llama.cpp

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

Sigbjørn Skjæretggml : allow fill node alloc inplace (#17870)86a3f0fad

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
1.1 KiB37 linesraw
1#include "fill.cuh"
2#include "convert.cuh"
3
4#define CUDA_FILL_BLOCK_SIZE 256
5
6template <typename T>
7static __global__ void fill_kernel(T * dst, const int64_t k, const T value) {
8    const int64_t i = (int64_t)blockDim.x * blockIdx.x + threadIdx.x;
9    if (i >= k) {
10        return;
11    }
12    dst[i] = value;
13}
14
15void ggml_cuda_op_fill(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
16    void * dst_d = dst->data;
17    cudaStream_t stream = ctx.stream();
18
19    GGML_ASSERT(ggml_is_contiguous(dst));
20
21    float value;
22    memcpy(&value, dst->op_params, sizeof(float));
23
24    const int64_t k = ggml_nelements(dst);
25    const int64_t num_blocks = (k + CUDA_FILL_BLOCK_SIZE - 1) / CUDA_FILL_BLOCK_SIZE;
26
27    switch (dst->type) {
28        case GGML_TYPE_F32:
29            fill_kernel<<<num_blocks, CUDA_FILL_BLOCK_SIZE, 0, stream>>>((float *)dst_d, k, value);
30            break;
31        case GGML_TYPE_F16:
32            fill_kernel<<<num_blocks, CUDA_FILL_BLOCK_SIZE, 0, stream>>>((half *)dst_d, k, ggml_cuda_cast<half>(value));
33            break;
34        default:
35            GGML_ABORT("unsupported type");
36    }
37}