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
brancheschar/aliceai-t5-moeall refs →
8.1 KiB249 linesraw
1#include "amx.h"
2#include "common.h"
3#include "mmq.h"
4#include "ggml-backend-impl.h"
5#include "ggml-backend.h"
6#include "ggml-impl.h"
7#include "ggml-cpu.h"
8#include "traits.h"
9
10#if defined(__linux__)
11#include <sys/syscall.h>
12#include <unistd.h>
13#endif
14
15#include <cstdlib>
16#include <cstring>
17#include <memory>
18
19#if defined(__AMX_INT8__) && defined(__AVX512VNNI__)
20
21// AMX type_trais
22namespace ggml::cpu::amx {
23class tensor_traits : public ggml::cpu::tensor_traits {
24    bool work_size(int /* n_threads */, const struct ggml_tensor * op, size_t & size) override {
25        size = ggml_backend_amx_desired_wsize(op);
26        return true;
27    }
28
29    bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) override {
30        if (op->op == GGML_OP_MUL_MAT) {
31            ggml_backend_amx_mul_mat(params, op);
32            return true;
33        }
34        return false;
35    }
36};
37
38static ggml::cpu::tensor_traits * get_tensor_traits(ggml_backend_buffer_t, struct ggml_tensor *) {
39    static tensor_traits traits;
40    return &traits;
41}
42}  // namespace ggml::cpu::amx
43
44// AMX buffer interface
45static void ggml_backend_amx_buffer_free_buffer(ggml_backend_buffer_t buffer) {
46    free(buffer->context);
47}
48
49static void * ggml_backend_amx_buffer_get_base(ggml_backend_buffer_t buffer) {
50    return (void *) (buffer->context);
51}
52
53static enum ggml_status ggml_backend_amx_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) {
54    tensor->extra = (void *) ggml::cpu::amx::get_tensor_traits(buffer, tensor);
55
56    GGML_UNUSED(buffer);
57    return GGML_STATUS_SUCCESS;
58}
59
60static void ggml_backend_amx_buffer_memset_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor,
61                                                  uint8_t value, size_t offset, size_t size) {
62    memset((char *) tensor->data + offset, value, size);
63
64    GGML_UNUSED(buffer);
65}
66
67static void ggml_backend_amx_buffer_set_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor,
68                                               const void * data, size_t offset, size_t size) {
69    if (qtype_has_amx_kernels(tensor->type)) {
70        GGML_LOG_DEBUG("%s: amx repack tensor %s of type %s\n", __func__, tensor->name, ggml_type_name(tensor->type));
71        ggml_backend_amx_convert_weight(tensor, data, offset, size);
72    } else {
73        memcpy((char *) tensor->data + offset, data, size);
74    }
75
76    GGML_UNUSED(buffer);
77}
78
79/*
80// need to figure what we need to do with buffer->extra.
81static void ggml_backend_amx_buffer_get_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) {
82    GGML_ASSERT(!qtype_has_amx_kernels(tensor->type));
83    memcpy(data, (const char *)tensor->data + offset, size);
84
85    GGML_UNUSED(buffer);
86}
87
88static bool ggml_backend_amx_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * src, struct ggml_tensor * dst) {
89    if (ggml_backend_buffer_is_host(src->buffer)) {
90        if (qtype_has_amx_kernels(src->type)) {
91            ggml_backend_amx_convert_weight(dst, src->data, 0, ggml_nbytes(dst));
92        } else {
93            memcpy(dst->data, src->data, ggml_nbytes(src));
94        }
95        return true;
96    }
97    return false;
98
99    GGML_UNUSED(buffer);
100}
101*/
102
103static void ggml_backend_amx_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
104    memset(buffer->context, value, buffer->size);
105}
106
107static ggml_backend_buffer_i ggml_backend_amx_buffer_interface = {
108    /* .free_buffer     = */ ggml_backend_amx_buffer_free_buffer,
109    /* .get_base        = */ ggml_backend_amx_buffer_get_base,
110    /* .init_tensor     = */ ggml_backend_amx_buffer_init_tensor,
111    /* .memset_tensor   = */ ggml_backend_amx_buffer_memset_tensor,
112    /* .set_tensor      = */ ggml_backend_amx_buffer_set_tensor,
113    /* .get_tensor      = */ nullptr,
114    /* .set_tensor_2d   = */ nullptr,
115    /* .get_tensor_2d   = */ nullptr,
116    /* .cpy_tensor      = */ nullptr,
117    /* .clear           = */ ggml_backend_amx_buffer_clear,
118    /* .reset           = */ nullptr,
119};
120
121static const char * ggml_backend_amx_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
122    return "AMX";
123
124    GGML_UNUSED(buft);
125}
126
127static ggml_backend_buffer_t ggml_backend_amx_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
128    void * data = ggml_aligned_malloc(size);
129    if (data == NULL) {
130        fprintf(stderr, "%s: failed to allocate buffer of size %zu\n", __func__, size);
131        return NULL;
132    }
133
134    return ggml_backend_buffer_init(buft, ggml_backend_amx_buffer_interface, data, size);
135}
136
137static size_t ggml_backend_amx_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) {
138    return TENSOR_ALIGNMENT;
139
140    GGML_UNUSED(buft);
141}
142
143namespace ggml::cpu::amx {
144class extra_buffer_type : ggml::cpu::extra_buffer_type {
145    bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override {
146        if (op->op != GGML_OP_MUL_MAT) {
147            return false;
148        }
149        auto * src0 = op->src[0];
150        auto * src1 = op->src[1];
151
152        if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1)) {
153            return false;
154        }
155        if (!src0->buffer || src0->buffer->buft != ggml_backend_amx_buffer_type()) {
156            return false;
157        }
158        if (src1->buffer && !ggml_backend_buft_is_host(src1->buffer->buft)) {
159            return false;
160        }
161        if (op->ne[0] % (TILE_N * 2)) {
162            return false;
163        }
164        int alignment;
165        switch (src0->type) {
166            case GGML_TYPE_Q4_0:
167            case GGML_TYPE_Q4_1:
168            case GGML_TYPE_Q8_0:
169                alignment = TILE_K;
170                break;
171            case GGML_TYPE_Q4_K:
172            case GGML_TYPE_Q5_K:
173            case GGML_TYPE_Q6_K:
174            case GGML_TYPE_IQ4_XS:
175                alignment = 256; // QK_K
176                break;
177            case GGML_TYPE_F16:
178                alignment = 16;
179                break;
180            default:
181                return false;
182        }
183        if (src0->ne[0] % alignment) {
184            return false;
185        }
186        if (src1->type != GGML_TYPE_F32) {
187            return false;
188        }
189        return true;
190    }
191
192    ggml::cpu::tensor_traits * get_tensor_traits(const struct ggml_tensor * op) override {
193        if (op->op == GGML_OP_MUL_MAT && op->src[0]->buffer &&
194            op->src[0]->buffer->buft == ggml_backend_amx_buffer_type()) {
195            return (ggml::cpu::tensor_traits *) op->src[0]->extra;
196        }
197
198        return nullptr;
199    }
200};
201}  // namespace ggml::cpu::amx
202
203static size_t ggml_backend_amx_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
204    return ggml_backend_amx_get_alloc_size(tensor);
205
206    GGML_UNUSED(buft);
207}
208
209#define ARCH_GET_XCOMP_PERM     0x1022
210#define ARCH_REQ_XCOMP_PERM     0x1023
211#define XFEATURE_XTILECFG       17
212#define XFEATURE_XTILEDATA      18
213
214static bool ggml_amx_init() {
215#if defined(__linux__)
216    if (syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA)) {
217        fprintf(stderr, "AMX is not ready to be used!\n");
218        return false;
219    }
220    return true;
221#elif defined(_WIN32)
222    return true;
223#else
224    return false;
225#endif
226}
227
228ggml_backend_buffer_type_t ggml_backend_amx_buffer_type() {
229    static struct ggml_backend_buffer_type ggml_backend_buffer_type_amx = {
230        /* .iface = */ {
231                        /* .get_name         = */ ggml_backend_amx_buffer_type_get_name,
232                        /* .alloc_buffer     = */ ggml_backend_amx_buffer_type_alloc_buffer,
233                        /* .get_alignment    = */ ggml_backend_amx_buffer_type_get_alignment,
234                        /* .get_max_size     = */ nullptr,  // defaults to SIZE_MAX
235                        /* .get_alloc_size   = */ ggml_backend_amx_buffer_type_get_alloc_size,
236                        /* .is_host          = */ nullptr,
237                        },
238        /* .device  = */ ggml_backend_reg_dev_get(ggml_backend_cpu_reg(), 0),
239        /* .context = */ new ggml::cpu::amx::extra_buffer_type(),
240    };
241
242    if (!ggml_amx_init()) {
243        return nullptr;
244    }
245
246    return &ggml_backend_buffer_type_amx;
247}
248
249#endif  // defined(__AMX_INT8__) && defined(__AVX512VNNI__)