char-fork/llama.cpp

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

Max Krasnyanskyhexagon: support for multi-device model split (aka row-split) (#28589)eafe15a5e

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
13.1 KiB328 linesraw
1#include "hex-common.h"
2#include "hex-profile.h"
3#include "htp-ctx.h"
4#include "htp-ops.h"
5#include "htp-tensor.h"
6#include "hexagon_types.h"
7#include "hexagon_protos.h"
8#include "hvx_hexagon_protos.h"
9#include "hex-dma.h"
10#include "htp-vtcm.h"
11#include "hvx-utils.h"
12#include "hex-fastdiv.h"
13#include <string.h>
14
15struct htp_concat_context {
16    struct htp_ops_context * octx;
17    uint32_t dim;
18    uint32_t nrows_per_thread;
19    uint32_t row_start;
20    uint32_t nrows;
21    uint32_t elem_start;
22    uint32_t nelems;
23    struct fastdiv_values div_ne0;
24    struct fastdiv_values div_ne1;
25    struct fastdiv_values div_ne2;
26};
27
28static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void * data) {
29    struct htp_concat_context * cctx = (struct htp_concat_context *) data;
30    struct htp_ops_context * octx = cctx->octx;
31
32    const struct htp_tensor * src0 = octx->src[0];
33    const struct htp_tensor * src1 = octx->src[1];
34    const struct htp_tensor * dst  = octx->dst;
35
36    const uint32_t src0_ne0 = src0->ne[0];
37    const uint32_t src1_ne0 = src1->ne[0];
38
39    const uint32_t row_end = cctx->row_start + cctx->nrows;
40    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
41    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
42    if (start_i >= end_i) return;
43
44    dma_queue * q = octx->ctx->dma[ith];
45
46    uint8_t * spad0_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread;
47    uint8_t * spad1_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread;
48
49    const uint32_t block_i = 32;
50    const uint32_t spad1_stride = block_i * sizeof(float);
51
52    int32_t offsets[32] __attribute__((aligned(128)));
53    for(int k=0; k<32; k++) {
54        offsets[k] = k * spad1_stride;
55    }
56    HVX_Vector vv = *(HVX_Vector*)offsets;
57    const uint32_t src1_ne0_padded = hex_round_up(src1_ne0, 32);
58    const uint32_t spad0_row_bytes = hex_round_up((src0_ne0 + src1_ne0_padded) * sizeof(float), VLEN);
59    uint32_t mu = src1_ne0_padded * spad1_stride;
60
61    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
62
63    for (uint32_t i = start_i; i < end_i; i += block_i) {
64        uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;
65
66        uint32_t src1_width_bytes = current_block_i * sizeof(float);
67        uint8_t * src1_ptr = (uint8_t *)src1->data + i * src1->nb[1];
68        dma_queue_push(q, dma_make_ptr(spad1_base, src1_ptr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0);
69
70        uint32_t src0_row_bytes = src0_ne0 * sizeof(float);
71        uint8_t * src0_ptr = (uint8_t *)src0->data + i * src0->nb[1];
72        dma_queue_push(q, dma_make_ptr(spad0_base, src0_ptr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i);
73
74        dma_queue_pop(q); // src1
75
76        HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);
77
78        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
79        for (uint32_t j = 0; j < src1_ne0_padded; j += 32) {
80            #pragma unroll(4)
81            for (uint32_t ii = 0; ii < current_block_i; ii++) {
82                size_t rt = (size_t)(spad1_base + j * spad1_stride + ii * sizeof(float));
83                Q6_vgather_ARMVw(&vtcm_tmp[ii], rt, mu, vv);
84                uint8_t * dst_ptr = spad0_base + ii * spad0_row_bytes + (src0_ne0 + j) * sizeof(float);
85                hvx_vmemu(dst_ptr) = vtcm_tmp[ii];
86            }
87        }
88        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
89
90        dma_queue_pop(q); // src0
91
92        uint8_t * dst_ptr = (uint8_t *)dst->data + i * dst->nb[1];
93        dma_queue_push(q, dma_make_ptr(dst_ptr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i);
94
95        dma_queue_pop(q);
96    }
97}
98
99static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void * data) {
100    struct htp_concat_context * cctx = (struct htp_concat_context *) data;
101    struct htp_ops_context * octx = cctx->octx;
102
103    const struct htp_tensor * src0 = octx->src[0];
104    const struct htp_tensor * src1 = octx->src[1];
105    const struct htp_tensor * dst  = octx->dst;
106
107    const uint32_t src0_ne0 = src0->ne[0];
108    const uint32_t src1_ne0 = src1->ne[0];
109
110    const uint32_t row_end = cctx->row_start + cctx->nrows;
111    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
112    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
113    if (start_i >= end_i) return;
114
115    dma_queue * q = octx->ctx->dma[ith];
116
117    uint8_t * spad0_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread;
118    uint8_t * spad1_base = octx->src1_spad.data + ith * octx->src1_spad.size_per_thread;
119
120    const uint32_t block_i = 64;
121    const uint32_t spad1_stride = block_i * sizeof(__fp16);
122
123    int16_t offsets[64] __attribute__((aligned(128)));
124    for(int k=0; k<64; k++) {
125        offsets[k] = k * spad1_stride;
126    }
127    HVX_Vector vv = *(HVX_Vector*)offsets;
128    const uint32_t src1_ne0_padded = hex_round_up(src1_ne0, 64);
129    const uint32_t spad0_row_bytes = hex_round_up((src0_ne0 + src1_ne0_padded) * sizeof(__fp16), VLEN);
130    uint32_t mu = src1_ne0_padded * spad1_stride;
131
132    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
133
134    for (uint32_t i = start_i; i < end_i; i += block_i) {
135        uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;
136
137        uint32_t src1_width_bytes = current_block_i * sizeof(__fp16);
138        uint8_t * src1_ptr = (uint8_t *)src1->data + i * src1->nb[1];
139        dma_queue_push(q, dma_make_ptr(spad1_base, src1_ptr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0);
140
141        uint32_t src0_row_bytes = src0_ne0 * sizeof(__fp16);
142        uint8_t * src0_ptr = (uint8_t *)src0->data + i * src0->nb[1];
143        dma_queue_push(q, dma_make_ptr(spad0_base, src0_ptr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i);
144
145        dma_queue_pop(q); // src1
146
147        HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);
148
149        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
150        for (uint32_t j = 0; j < src1_ne0_padded; j += 64) {
151            #pragma unroll(4)
152            for (uint32_t ii = 0; ii < current_block_i; ii++) {
153                size_t rt = (size_t)(spad1_base + j * spad1_stride + ii * sizeof(__fp16));
154                Q6_vgather_ARMVh(&vtcm_tmp[ii], rt, mu, vv);
155                uint8_t * dst_ptr = spad0_base + ii * spad0_row_bytes + (src0_ne0 + j) * sizeof(__fp16);
156                hvx_vmemu(dst_ptr) = vtcm_tmp[ii];
157            }
158        }
159        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
160
161        dma_queue_pop(q); // src0
162
163        uint8_t * dst_ptr = (uint8_t *)dst->data + i * dst->nb[1];
164        dma_queue_push(q, dma_make_ptr(dst_ptr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i);
165
166        dma_queue_pop(q);
167    }
168}
169
170static void concat_generic(unsigned int nth, unsigned int ith, void * data) {
171    struct htp_concat_context * cctx = (struct htp_concat_context *) data;
172    struct htp_ops_context * octx = cctx->octx;
173
174    const struct htp_tensor * src0 = octx->src[0];
175    const struct htp_tensor * src1 = octx->src[1];
176    const struct htp_tensor * dst  = octx->dst;
177
178    const int dim = cctx->dim;
179    const uint32_t type_size = (dst->type == HTP_TYPE_F32 || dst->type == HTP_TYPE_I32) ? 4 : 2;
180
181    const uint32_t ne[4] = {dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]};
182
183    // Per-device element range aligned to prevent false sharing
184    const uint32_t elem_start = cctx->elem_start;
185    const uint32_t nelems     = cctx->nelems;
186    const uint32_t chunk_size = (nelems + nth - 1) / nth;
187
188    const uint32_t start_idx = MIN(elem_start + ith * chunk_size, elem_start + nelems);
189    const uint32_t end_idx   = MIN(start_idx + chunk_size, elem_start + nelems);
190
191    // Naive scalar element-wise copy
192    for (uint32_t idx = start_idx; idx < end_idx; idx++) {
193        uint32_t idx_div_ne0 = fastdiv(idx, &cctx->div_ne0);
194        uint32_t i0 = idx - idx_div_ne0 * ne[0];
195
196        uint32_t idx_div_ne01 = fastdiv(idx_div_ne0, &cctx->div_ne1);
197        uint32_t i1 = idx_div_ne0 - idx_div_ne01 * ne[1];
198
199        uint32_t idx_div_ne012 = fastdiv(idx_div_ne01, &cctx->div_ne2);
200        uint32_t i2 = idx_div_ne01 - idx_div_ne012 * ne[2];
201        uint32_t i3 = idx_div_ne012;
202
203        uint8_t * dst_ptr = (uint8_t *)dst->data + i3 * dst->nb[3] + i2 * dst->nb[2] + i1 * dst->nb[1] + i0 * dst->nb[0];
204
205        uint32_t idx_dim = 0;
206        if (dim == 0) idx_dim = i0;
207        else if (dim == 1) idx_dim = i1;
208        else if (dim == 2) idx_dim = i2;
209        else if (dim == 3) idx_dim = i3;
210
211        const struct htp_tensor * src = (idx_dim < src0->ne[dim]) ? src0 : src1;
212
213        uint32_t s0 = i0;
214        uint32_t s1 = i1;
215        uint32_t s2 = i2;
216        uint32_t s3 = i3;
217
218        if (dim == 0 && src == src1) s0 -= src0->ne[0];
219        if (dim == 1 && src == src1) s1 -= src0->ne[1];
220        if (dim == 2 && src == src1) s2 -= src0->ne[2];
221        if (dim == 3 && src == src1) s3 -= src0->ne[3];
222
223        uint8_t * src_ptr = (uint8_t *)src->data + s3 * src->nb[3] + s2 * src->nb[2] + s1 * src->nb[1] + s0 * src->nb[0];
224
225        if (type_size == 4) {
226            *(float*)dst_ptr = *(float*)src_ptr;
227        } else {
228            *(__fp16*)dst_ptr = *(__fp16*)src_ptr;
229        }
230    }
231}
232
233int op_concat(struct htp_ops_context * octx) {
234    const struct htp_tensor * src0 = octx->src[0];
235    const struct htp_tensor * src1 = octx->src[1];
236    const struct htp_tensor * dst  = octx->dst;
237
238    int dim = octx->op_params[0];
239
240    bool is_2d = dst->ne[2] == 1 && dst->ne[3] == 1;
241
242    const uint32_t type_size = (dst->type == HTP_TYPE_F32 || dst->type == HTP_TYPE_I32) ? 4 : 2;
243    bool is_src1_transposed  = (src1->nb[0] > src1->nb[1]);
244    bool is_src0_transposed  = (src0->nb[0] > src0->nb[1]);
245
246    uint32_t n_threads = octx->n_threads;
247    struct htp_concat_context cctx;
248    cctx.octx = octx;
249    cctx.dim = dim;
250    cctx.div_ne0 = init_fastdiv_values(dst->ne[0]);
251    cctx.div_ne1 = init_fastdiv_values(dst->ne[1]);
252    cctx.div_ne2 = init_fastdiv_values(dst->ne[2]);
253
254    void (*worker_func)(unsigned int, unsigned int, void *) = concat_generic;
255
256    if (dim == 0 && is_2d && is_src1_transposed && !is_src0_transposed) {
257        const uint32_t total_rows = dst->ne[1];
258        const size_t dst_data_row_size = dst->ne[0] * type_size;
259        uint32_t row_start = 0;
260        uint32_t nrows     = total_rows;
261        if (octx->ctx->mdev.count > 1) {
262            uint32_t rows_per_chunk = 0;
263            htp_tensor_mdev_rows_per_chunk(dst, type_size, (uint32_t) dst_data_row_size, &rows_per_chunk);
264            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
265            row_start = range.start;
266            nrows     = range.count;
267        }
268
269        if (nrows == 0) {
270            return HTP_STATUS_OK;
271        }
272
273        cctx.row_start = row_start;
274        cctx.nrows     = nrows;
275
276        uint32_t block_i = (type_size == 4) ? 32 : 64;
277
278        cctx.nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
279
280        // Allocate VTCM
281        uint32_t spad1_stride = block_i * type_size;
282
283        uint32_t src1_ne0_padded = hex_round_up(src1->ne[0], block_i);
284        uint32_t spad0_row_bytes = hex_round_up((src0->ne[0] + src1_ne0_padded) * type_size, VLEN);
285
286        octx->src0_spad.size_per_thread = block_i * spad0_row_bytes;
287        octx->src1_spad.size_per_thread = src1_ne0_padded * spad1_stride + block_i * VLEN;
288
289        octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
290        octx->src1_spad.size = n_threads * octx->src1_spad.size_per_thread;
291
292        if (octx->src0_spad.size + octx->src1_spad.size > octx->ctx->vtcm_size) {
293            return HTP_STATUS_VTCM_TOO_SMALL;
294        }
295
296        octx->src0_spad.data = octx->ctx->vtcm_base;
297        octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size;
298        octx->src0_spad.src  = NULL;
299        octx->src1_spad.src  = NULL;
300
301        if (type_size == 4) {
302            worker_func = concat_2d_f32_transposed;
303        } else {
304            worker_func = concat_2d_f16_transposed;
305        }
306    } else {
307        const uint32_t total_elements = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
308        uint32_t elem_start = 0;
309        uint32_t nelems     = total_elements;
310        if (octx->ctx->mdev.count > 1) {
311            const uint32_t elems_per_chunk = HEX_L2_LINE_SIZE / type_size;
312            const bool can_split = htp_tensor_mdev_data_aligned(dst) && htp_tensor_is_contiguous(dst, type_size) && !htp_tensor_is_permuted(dst);
313            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_elements, can_split ? elems_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
314            elem_start = range.start;
315            nelems     = range.count;
316        }
317
318        if (nelems == 0) {
319            return HTP_STATUS_OK;
320        }
321
322        cctx.elem_start = elem_start;
323        cctx.nelems     = nelems;
324    }
325
326    work_queue_run(octx->ctx->work_queue, worker_func, &cctx, n_threads);
327    return HTP_STATUS_OK;
328}