char-fork/llama.cpp

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

Johannes Gäßlerggml: backend-agnostic tensor parallelism (experimental) (#19378)d6f303004

char/aliceai-t5-moe
5.3 KiB123 linesraw
1#include "ggml-remoting.h"
2
3#define BUFFER_TO_GPU(name) ((ggml_backend_remoting_buffer_context *) (name)->context)->gpu
4
5static void * ggml_backend_remoting_buffer_get_base(ggml_backend_buffer_t buffer) {
6    ggml_backend_remoting_buffer_context * context = (ggml_backend_remoting_buffer_context *) buffer->context;
7    if (context->base) {
8        return context->base;
9    }
10
11    context->base = apir_buffer_get_base(BUFFER_TO_GPU(buffer), BUFFER_TO_APIR_CONTEXT(buffer));
12
13    return context->base;
14}
15
16static void ggml_backend_remoting_buffer_set_tensor(ggml_backend_buffer_t buffer,
17                                                    ggml_tensor *         tensor,
18                                                    const void *          data,
19                                                    size_t                offset,
20                                                    size_t                size) {
21    virtgpu * gpu = BUFFER_TO_GPU(buffer);
22
23    ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer);
24    if (context->is_from_ptr) {
25        memcpy((char *) tensor->data + offset, data, size);
26    } else {
27        apir_buffer_set_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), tensor, data, offset, size);
28    }
29
30    return;
31}
32
33static void ggml_backend_remoting_buffer_get_tensor(ggml_backend_buffer_t buffer,
34                                                    const ggml_tensor *   tensor,
35                                                    void *                data,
36                                                    size_t                offset,
37                                                    size_t                size) {
38    virtgpu *                              gpu     = BUFFER_TO_GPU(buffer);
39    ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer);
40    if (context->is_from_ptr) {
41        memcpy(data, (const char *) tensor->data + offset, size);
42    } else {
43        apir_buffer_get_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), tensor, data, offset, size);
44    }
45}
46
47static void ggml_backend_remoting_buffer_set_tensor_from_ptr(ggml_backend_buffer_t buffer,
48                                                             ggml_tensor *         tensor,
49                                                             const void *          data,
50                                                             size_t                offset,
51                                                             size_t                size) {
52    UNUSED(buffer);
53
54    memcpy((char *) tensor->data + offset, data, size);
55
56    return;
57}
58
59static void ggml_backend_remoting_buffer_get_tensor_from_ptr(ggml_backend_buffer_t buffer,
60                                                             const ggml_tensor *   tensor,
61                                                             void *                data,
62                                                             size_t                offset,
63                                                             size_t                size) {
64    UNUSED(buffer);
65
66    memcpy(data, (const char *) tensor->data + offset, size);
67}
68
69static bool ggml_backend_remoting_buffer_cpy_tensor(ggml_backend_buffer_t buffer,
70                                                    const ggml_tensor *   src,
71                                                    ggml_tensor *         dst) {
72    virtgpu * gpu = BUFFER_TO_GPU(buffer);
73
74    bool ret = apir_buffer_cpy_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), src, dst);
75
76    return ret;
77}
78
79static void ggml_backend_remoting_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
80    virtgpu * gpu = BUFFER_TO_GPU(buffer);
81
82    apir_buffer_clear(gpu, BUFFER_TO_APIR_CONTEXT(buffer), value);
83
84    return;
85}
86
87static void ggml_backend_remoting_buffer_free_buffer(ggml_backend_buffer_t buffer) {
88    virtgpu * gpu = BUFFER_TO_GPU(buffer);
89
90    apir_buffer_free_buffer(gpu, BUFFER_TO_APIR_CONTEXT(buffer));
91
92    ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer);
93    free(context);
94    buffer->context = NULL;
95}
96
97const ggml_backend_buffer_i ggml_backend_remoting_buffer_interface = {
98    /* .free_buffer     = */ ggml_backend_remoting_buffer_free_buffer,
99    /* .get_base        = */ ggml_backend_remoting_buffer_get_base,
100    /* .init_tensor     = */ NULL,
101    /* .memset_tensor   = */ NULL,
102    /* .set_tensor      = */ ggml_backend_remoting_buffer_set_tensor,
103    /* .get_tensor      = */ ggml_backend_remoting_buffer_get_tensor,
104    /* .set_tensor_2d   = */ NULL,
105    /* .get_tensor_2d   = */ NULL,
106    /* .cpy_tensor      = */ ggml_backend_remoting_buffer_cpy_tensor,
107    /* .clear           = */ ggml_backend_remoting_buffer_clear,
108    /* .reset           = */ NULL,
109};
110
111const ggml_backend_buffer_i ggml_backend_remoting_buffer_from_ptr_interface = {
112    /* .free_buffer     = */ ggml_backend_remoting_buffer_free_buffer,
113    /* .get_base        = */ ggml_backend_remoting_buffer_get_base,
114    /* .init_tensor     = */ NULL,
115    /* .memset_tensor   = */ NULL,
116    /* .set_tensor      = */ ggml_backend_remoting_buffer_set_tensor_from_ptr,
117    /* .get_tensor      = */ ggml_backend_remoting_buffer_get_tensor_from_ptr,
118    /* .set_tensor_2d   = */ NULL,
119    /* .get_tensor_2d   = */ NULL,
120    /* .cpy_tensor      = */ ggml_backend_remoting_buffer_cpy_tensor,
121    /* .clear           = */ ggml_backend_remoting_buffer_clear,
122    /* .reset           = */ NULL,
123};