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 →
5.6 KiB169 linesraw
1#pragma clang diagnostic ignored "-Wunused-variable"
2#pragma clang diagnostic ignored "-Wunused-function"
3#pragma clang diagnostic ignored "-Wunused-but-set-variable"
4
5#include <HAP_farf.h>
6#include <HAP_perf.h>
7
8#include <string.h>
9
10#include "hvx-utils.h"
11
12#define GGML_COMMON_DECL_C
13#include "ggml-common.h"
14#include "htp-ctx.h"
15#include "hex-common.h"
16#include "hex-profile.h"
17#include "htp-ops.h"
18#include "htp-tensor.h"
19
20struct htp_repeat_context {
21    struct htp_ops_context * octx;
22
23    uint32_t nr0;
24    uint32_t nr1;
25    uint32_t nr2;
26    uint32_t nr3;
27
28    uint32_t nrows_per_thread;
29    uint32_t total_dst_rows;  // ne1 * ne2 * ne3
30    uint32_t row_start;
31
32    size_t   type_size;
33};
34
35static void repeat_job_per_thread(unsigned int nth, unsigned int ith, void * data) {
36    const struct htp_repeat_context * rctx = (const struct htp_repeat_context *) data;
37    struct htp_ops_context * octx = rctx->octx;
38    const struct htp_tensor * src = octx->src[0];
39    const struct htp_tensor * dst = octx->dst;
40
41    const uint32_t ne00 = src->ne[0];
42    const uint32_t ne01 = src->ne[1];
43    const uint32_t ne02 = src->ne[2];
44    const uint32_t ne03 = src->ne[3];
45
46    const uint32_t nb00 = src->nb[0];
47    const uint32_t nb01 = src->nb[1];
48    const uint32_t nb02 = src->nb[2];
49    const uint32_t nb03 = src->nb[3];
50
51    const uint32_t ne0 = dst->ne[0];
52    const uint32_t ne1 = dst->ne[1];
53    const uint32_t ne2 = dst->ne[2];
54    const uint32_t ne3 = dst->ne[3];
55
56    const uint32_t nb0 = dst->nb[0];
57    const uint32_t nb1 = dst->nb[1];
58    const uint32_t nb2 = dst->nb[2];
59    const uint32_t nb3 = dst->nb[3];
60
61    const uint32_t nr0 = rctx->nr0;
62    const uint32_t nr1 = rctx->nr1;
63    const uint32_t nr2 = rctx->nr2;
64    const uint32_t nr3 = rctx->nr3;
65
66    const size_t row_bytes = ne00 * rctx->type_size;
67
68    const uint32_t row_start = rctx->row_start + rctx->nrows_per_thread * ith;
69    const uint32_t row_end   = MIN(row_start + rctx->nrows_per_thread, rctx->row_start + rctx->total_dst_rows);
70
71    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
72    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
73
74    for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
75        // Decompose flat dst row index into (i1, i2, i3)
76        const uint32_t i1 = dst_row % ne1;
77        const uint32_t i2 = (dst_row / ne1) % ne2;
78        const uint32_t i3 = dst_row / (ne1 * ne2);
79
80        // Map to source indices (tiling)
81        const uint32_t k1 = i1 % ne01;
82        const uint32_t k2 = i2 % ne02;
83        const uint32_t k3 = i3 % ne03;
84
85        const uint8_t * src_row = (const uint8_t *) src->data + k1 * nb01 + k2 * nb02 + k3 * nb03;
86        uint8_t * dst_base      = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
87
88        // Tile along dimension 0
89        for (uint32_t i0 = 0; i0 < nr0; i0++) {
90            uint8_t * dst_ptr = dst_base + i0 * ne00 * nb0;
91            memcpy(dst_ptr, src_row, row_bytes);
92        }
93    }
94
95    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
96
97    FARF(HIGH, "repeat %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
98         ith, nth, src->ne[0], src->ne[1], src->ne[2], src->ne[3],
99         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
100         row_start, row_end);
101}
102
103int op_repeat(struct htp_ops_context * octx) {
104    const struct htp_tensor * src0 = octx->src[0];
105    const struct htp_tensor * dst  = octx->dst;
106
107    // Validate that dst dims are multiples of src dims
108    if (dst->ne[0] % src0->ne[0] != 0 ||
109        dst->ne[1] % src0->ne[1] != 0 ||
110        dst->ne[2] % src0->ne[2] != 0 ||
111        dst->ne[3] % src0->ne[3] != 0) {
112        FARF(ERROR, "repeat: dst dims must be multiples of src dims\n");
113        return HTP_STATUS_INVAL_PARAMS;
114    }
115
116    size_t type_size;
117    switch (src0->type) {
118        case HTP_TYPE_F32: type_size = 4; break;
119        case HTP_TYPE_F16: type_size = 2; break;
120        default:
121            FARF(ERROR, "repeat: unsupported type %u\n", src0->type);
122            return HTP_STATUS_NO_SUPPORT;
123    }
124
125    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
126        return HTP_STATUS_OK;
127    }
128
129    const uint32_t total_dst_rows  = dst->ne[1] * dst->ne[2] * dst->ne[3];
130    const size_t dst_row_size = dst->ne[0] * type_size;
131
132    uint32_t row_start = 0;
133    uint32_t nrows     = total_dst_rows;
134
135    if (octx->ctx->mdev.count > 1) {
136        uint32_t rows_per_chunk = 0;
137        htp_tensor_mdev_rows_per_chunk(dst, type_size, (uint32_t) dst_row_size, &rows_per_chunk);
138        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_dst_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
139        row_start = range.start;
140        nrows     = range.count;
141    }
142
143    if (nrows == 0) {
144        return HTP_STATUS_OK;
145    }
146
147    const uint32_t n_threads = octx->n_threads;
148
149    struct htp_repeat_context rctx = {
150        .octx             = octx,
151        .nr0              = dst->ne[0] / src0->ne[0],
152        .nr1              = dst->ne[1] / src0->ne[1],
153        .nr2              = dst->ne[2] / src0->ne[2],
154        .nr3              = dst->ne[3] / src0->ne[3],
155        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
156        .total_dst_rows   = nrows,
157        .row_start        = row_start,
158        .type_size        = type_size,
159    };
160
161    FARF(HIGH, "repeat: (%ux%ux%ux%u) -> (%ux%ux%ux%u) nr=(%u,%u,%u,%u)\n",
162         src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
163         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
164         rctx.nr0, rctx.nr1, rctx.nr2, rctx.nr3);
165
166    work_queue_run(octx->ctx->work_queue, repeat_job_per_thread, &rctx, n_threads);
167
168    return HTP_STATUS_OK;
169}