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.3 KiB154 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#include <math.h>
10
11#include "hex-dma.h"
12#include "hvx-utils.h"
13
14#define GGML_COMMON_DECL_C
15#include "ggml-common.h"
16#include "hex-common.h"
17#include "hex-profile.h"
18#include "htp-ctx.h"
19#include "htp-ops.h"
20#include "htp-tensor.h"
21
22#define sum_rows_preamble                         \
23    const struct htp_tensor *src0 = octx->src[0]; \
24    const struct htp_tensor *dst  = octx->dst;    \
25                                                  \
26    const uint32_t ne00 = src0->ne[0];            \
27    const uint32_t ne01 = src0->ne[1];            \
28    const uint32_t ne02 = src0->ne[2];            \
29    const uint32_t ne03 = src0->ne[3];            \
30                                                  \
31    const uint32_t nb00 = src0->nb[0];            \
32    const uint32_t nb01 = src0->nb[1];            \
33    const uint32_t nb02 = src0->nb[2];            \
34    const uint32_t nb03 = src0->nb[3];            \
35                                                  \
36    const uint32_t  ne0 = dst->ne[0];             \
37    const uint32_t  ne1 = dst->ne[1];             \
38    const uint32_t  ne2 = dst->ne[2];             \
39    const uint32_t  ne3 = dst->ne[3];             \
40                                                  \
41    const uint32_t  nb0 = dst->nb[0];             \
42    const uint32_t  nb1 = dst->nb[1];             \
43    const uint32_t  nb2 = dst->nb[2];             \
44    const uint32_t  nb3 = dst->nb[3];             \
45
46struct sum_rows_context {
47    struct htp_ops_context * octx;
48    const uint8_t * src_data;
49    uint8_t       * dst_data;
50    uint32_t        ne00;
51    size_t          src_stride;
52    size_t          dst_stride;
53    uint32_t        rows_per_thread;
54    uint32_t        total_rows;
55    bool            opt_path;
56};
57
58static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data) {
59    const struct sum_rows_context * smctx = (const struct sum_rows_context *) data;
60
61    const uint32_t rows_per_thread = smctx->rows_per_thread;
62    const uint32_t total_rows      = smctx->total_rows;
63
64    const uint32_t start_row = rows_per_thread * ith;
65    const uint32_t end_row   = MIN(start_row + rows_per_thread, total_rows);
66
67    if (start_row >= end_row) {
68        return;
69    }
70
71    const size_t   src_stride = smctx->src_stride;
72    const size_t   dst_stride = smctx->dst_stride;
73    const uint32_t ne00       = smctx->ne00;
74    const bool     opt_path   = smctx->opt_path;
75
76    const float * restrict src_th = (const float *) (smctx->src_data + (start_row * src_stride));
77    float       * restrict dst_th = (float *)       (smctx->dst_data + (start_row * dst_stride));
78
79    // Calculate actual number of rows for this thread
80    const uint32_t n_rows = end_row - start_row;
81
82    struct htp_thread_trace * tr = &smctx->octx->ctx->trace[ith];
83    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_row);
84
85    for (uint32_t ir = 0; ir < n_rows; ir++) {
86        const float * restrict src_local = src_th + (ir * (src_stride / sizeof(float)));
87
88        if (ir + 1 < n_rows) {
89            hex_l2fetch(src_local + (src_stride / sizeof(float)), src_stride, src_stride, 1);
90        }
91
92        if (opt_path) {
93            dst_th[ir] = hvx_reduce_sum_f32_a((const uint8_t *) src_local, ne00);
94        } else {
95            dst_th[ir] = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
96        }
97    }
98
99    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_row);
100}
101
102int op_sum_rows(struct htp_ops_context * octx) {
103    sum_rows_preamble;
104
105    if (octx->src[0]->type != HTP_TYPE_F32) {
106        return HTP_STATUS_NO_SUPPORT;
107    }
108
109    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
110        return HTP_STATUS_OK;
111    }
112
113    const uint32_t src0_nrows      = ne01 * ne02 * ne03;
114    const size_t dst_data_row_size = dst->ne[0] * sizeof(float);
115
116    uint32_t row_start = 0;
117    uint32_t nrows     = src0_nrows;
118
119    if (octx->ctx->mdev.count > 1) {
120        uint32_t rows_per_chunk = 0;
121        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_data_row_size, &rows_per_chunk);
122        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
123        row_start = range.start;
124        nrows     = range.count;
125    }
126
127    if (nrows == 0) {
128        return HTP_STATUS_OK;
129    }
130
131    const uint32_t n_threads = octx->n_threads;
132    const uint32_t rows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
133
134    bool opt_path = false;
135    if ((0 == hex_is_aligned((void *) src0->data, VLEN)) && !(nb01 & (VLEN - 1))) {
136        opt_path = true;
137    }
138
139    struct sum_rows_context smctx = {
140        .octx            = octx,
141        .src_data        = (const uint8_t *) src0->data + row_start * nb01,
142        .dst_data        = (uint8_t *) dst->data + row_start * nb1,
143        .ne00            = ne00,
144        .src_stride      = nb01,
145        .dst_stride      = nb1,
146        .rows_per_thread = rows_per_thread,
147        .total_rows      = nrows,
148        .opt_path        = opt_path,
149    };
150
151    work_queue_run(octx->ctx->work_queue, sum_rows_thread_f32, &smctx, n_threads);
152
153    return HTP_STATUS_OK;
154}