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 →
9.7 KiB247 linesraw
1#pragma clang diagnostic ignored "-Wunused-but-set-variable"
2
3#include <HAP_farf.h>
4#include <HAP_perf.h>
5
6#define GGML_COMMON_DECL_C
7#include "ggml-common.h"
8#include "hex-common.h"
9#include "hex-profile.h"
10#include "htp-ctx.h"
11#include "htp-ops.h"
12#include "htp-tensor.h"
13#include "hvx-types.h"
14#include "hex-utils.h"
15#include "hvx-copy.h"
16#include "hex-dma.h"
17
18#define htp_diag_tensors_preamble                           \
19    const struct htp_tensor * restrict src0 = octx->src[0]; \
20    const struct htp_tensor * restrict dst  = octx->dst;    \
21                                                            \
22    const uint32_t ne02 = src0->ne[2];                      \
23                                                            \
24    const uint32_t ne0 = dst->ne[0];                        \
25    const uint32_t ne1 = dst->ne[1];                        \
26                                                            \
27    const uint32_t nb02 = src0->nb[2];                      \
28    const uint32_t nb03 = src0->nb[3];                      \
29                                                            \
30    const uint32_t nb1 = dst->nb[1];                        \
31    const uint32_t nb2 = dst->nb[2];                        \
32    const uint32_t nb3 = dst->nb[3];
33
34struct htp_diag_context {
35    struct htp_ops_context * octx;
36    size_t          src_batch_size;
37    size_t          dst_row_size;
38    size_t          src_batch_size_aligned;
39    size_t          dst_row_size_aligned;
40    uint32_t        batches_per_thread;
41    uint32_t        total_batches;
42    uint32_t        batch_start;
43};
44
45#define htp_diag_preamble                                              \
46    struct htp_diag_context * dctx = (struct htp_diag_context *) data; \
47    struct htp_ops_context *  octx = dctx->octx;                       \
48    htp_diag_tensors_preamble;
49
50static inline void hvx_diag_row_f32(const float * restrict src, float * restrict dst,
51                                    uint32_t row_idx, uint32_t n) {
52    hvx_splat_f32_a((uint8_t *) dst, 0.0f, n);
53    dst[row_idx] = src[row_idx];
54}
55
56// ---------------------------------------------------------------------------
57// Per thread worker: DMA src fetch, compute in VTCM, DMA dst writeback
58// ---------------------------------------------------------------------------
59
60static void diag_thread_f32_dma(unsigned int nth, unsigned int ith, void * data) {
61    htp_diag_preamble;
62    dma_queue * dma_queue = octx->ctx->dma[ith];
63
64    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith;
65    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches);
66
67    if (ib0 >= ib1) {
68        return;
69    }
70
71    const size_t src_batch_size         = dctx->src_batch_size;
72    const size_t dst_row_size           = dctx->dst_row_size;
73    const size_t src_batch_size_aligned = dctx->src_batch_size_aligned;
74    const size_t dst_row_size_aligned   = dctx->dst_row_size_aligned;
75
76    const uint8_t * src_data = (const uint8_t *) src0->data;
77    uint8_t *       dst_data = (uint8_t *) dst->data;
78
79    // 1 src buffer + 1 dst row buffer per thread in VTCM
80    uint8_t * src_spad = octx->src0_spad.data + (ith * src_batch_size_aligned);
81    uint8_t * dst_spad = octx->dst_spad.data  + (ith * dst_row_size_aligned);
82
83    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
84
85    for (uint32_t ib = ib0; ib < ib1; ib++) {
86        const uint32_t i3 = ib / ne02;
87        const uint32_t i2 = ib % ne02;
88
89        const uint8_t * src_batch = src_data + i3 * nb03 + i2 * nb02;
90
91        // Fetch source vector into VTCM
92        dma_queue_push_ddr_to_vtcm(dma_queue,
93                                   dma_make_ptr(src_spad, src_batch),
94                                   src_batch_size_aligned, src_batch_size, 1);
95        dma_queue_flush(dma_queue);
96
97        const float * src_spad_f32 = (const float *) src_spad;
98        float       * dst_spad_f32 = (float *) dst_spad;
99
100        for (uint32_t i1 = 0; i1 < ne1; i1++) {
101            // Compute row in VTCM
102            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (ib * ne1 + i1));
103            hvx_diag_row_f32(src_spad_f32, dst_spad_f32, i1, ne0);
104            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (ib * ne1 + i1));
105
106            // Write completed row back to DDR
107            uint8_t * dst_row = dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1;
108            dma_queue_push_vtcm_to_ddr(dma_queue,
109                                       dma_make_ptr(dst_row, dst_spad),
110                                       dst_row_size, dst_row_size_aligned, 1);
111            dma_queue_flush(dma_queue);
112        }
113    }
114
115    FARF(HIGH, "diag-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
116         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1,
117         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
118}
119
120// ---------------------------------------------------------------------------
121// Per thread worker: Direct HVX (no DMA)
122// ---------------------------------------------------------------------------
123
124static void diag_thread_f32(unsigned int nth, unsigned int ith, void * data) {
125    htp_diag_preamble;
126
127    const uint8_t * src_data = (const uint8_t *) src0->data;
128    uint8_t *       dst_data = (uint8_t *) dst->data;
129
130    const uint32_t ib0 = dctx->batch_start + dctx->batches_per_thread * ith;
131    const uint32_t ib1 = MIN(ib0 + dctx->batches_per_thread, dctx->batch_start + dctx->total_batches);
132
133    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
134    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);
135
136    for (uint32_t ib = ib0; ib < ib1; ib++) {
137        const uint32_t i3 = ib / ne02;
138        const uint32_t i2 = ib % ne02;
139
140        const float * restrict src_batch = (const float *)(src_data + i3 * nb03 + i2 * nb02);
141
142        for (uint32_t i1 = 0; i1 < ne1; i1++) {
143            float * restrict dst_row = (float *)(dst_data + i3 * nb3 + i2 * nb2 + i1 * nb1);
144            hvx_diag_row_f32(src_batch, dst_row, i1, ne0);
145        }
146    }
147
148    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ib0);
149
150    FARF(HIGH, "diag-f32 %d/%d: %ux%ux%ux%u (%u:%u) -> %ux%ux%ux%u\n",
151         ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], ib0, ib1,
152         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]);
153}
154
155int op_diag_f32(struct htp_ops_context * octx) {
156    const struct htp_tensor * src0 = octx->src[0];
157    const struct htp_tensor * dst  = octx->dst;
158
159    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
160        return HTP_STATUS_OK;
161    }
162
163    const uint32_t total_batches = src0->ne[2] * src0->ne[3];
164    const size_t dst_batch_size  = dst->ne[1] * dst->nb[1];
165
166    uint32_t batch_start = 0;
167    uint32_t nbatches    = total_batches;
168
169    if (octx->ctx->mdev.count > 1) {
170        bool can_split = htp_tensor_mdev_data_aligned(dst) && (dst->ne[0] == 1 || dst->nb[0] == sizeof(float)) && !htp_tensor_is_permuted(dst);
171        uint32_t batches_per_chunk = 1;
172        if (can_split) {
173            if (dst->ne[2] > 1 && (dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
174                (dst->ne[3] <= 1 || (dst->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
175                batches_per_chunk = 1;
176            } else if (dst->nb[2] == dst_batch_size &&
177                       (dst->ne[3] <= 1 || dst->nb[3] == dst->nb[2] * dst->ne[2])) {
178                batches_per_chunk = (dst_batch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(dst_batch_size, HEX_L2_LINE_SIZE)) : 1;
179            } else {
180                can_split = false;
181            }
182        }
183
184        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_batches, can_split ? batches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
185        batch_start = range.start;
186        nbatches    = range.count;
187    }
188
189    if (nbatches == 0) {
190        return HTP_STATUS_OK;
191    }
192
193    const uint32_t n_threads     = octx->n_threads;
194
195    const size_t src_batch_size         = src0->ne[0] * sizeof(float);
196    const size_t dst_row_size           = dst->ne[0] * sizeof(float);
197    const size_t src_batch_size_aligned = hex_round_up(src_batch_size, VLEN);
198    const size_t dst_row_size_aligned   = hex_round_up(dst_row_size, VLEN);
199
200    // 1 src buffer + 1 dst row buffer per thread
201    const size_t spad_per_thread = src_batch_size_aligned + dst_row_size_aligned;
202
203    octx->src0_spad.size_per_thread = src_batch_size_aligned;
204    octx->dst_spad.size_per_thread  = dst_row_size_aligned;
205
206    octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
207    octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
208
209    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
210    octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size; octx->dst_spad.src  = NULL;
211
212    struct htp_diag_context dctx = {
213        .octx                   = octx,
214        .src_batch_size         = src_batch_size,
215        .dst_row_size           = dst_row_size,
216        .src_batch_size_aligned = src_batch_size_aligned,
217        .dst_row_size_aligned   = dst_row_size_aligned,
218        .batches_per_thread     = fastdiv(nbatches + n_threads - 1, &octx->n_threads_div),
219        .total_batches          = nbatches,
220        .batch_start            = batch_start,
221    };
222
223    if (octx->ctx->vtcm_size < spad_per_thread * n_threads) {
224        work_queue_run(octx->ctx->work_queue, diag_thread_f32, &dctx, n_threads);
225    } else {
226        work_queue_run(octx->ctx->work_queue, diag_thread_f32_dma, &dctx, n_threads);
227    }
228
229    return HTP_STATUS_OK;
230}
231
232int op_diag(struct htp_ops_context * octx) {
233    const struct htp_tensor * dst = octx->dst;
234
235    int err = HTP_STATUS_OK;
236
237    switch (dst->type) {
238        case HTP_TYPE_F32:
239            err = op_diag_f32(octx);
240            break;
241        default:
242            err = HTP_STATUS_NO_SUPPORT;
243            break;
244    }
245
246    return err;
247}