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 →
4.9 KiB146 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 <string.h>
7
8#include "hex-common.h"
9#include "hex-profile.h"
10
11#include "hvx-copy.h"
12#include "hvx-utils.h"
13
14#define GGML_COMMON_DECL_C
15#include "ggml-common.h"
16#include "htp-ctx.h"
17#include "htp-ops.h"
18#include "htp-tensor.h"
19
20// ggml op_params layout for FILL:
21//   op_params[0] (as float) - the scalar fill value
22
23#define fill_preamble                          \
24    const struct htp_tensor * dst = octx->dst; \
25                                               \
26    const uint32_t ne0 = dst->ne[0];           \
27    const uint32_t ne1 = dst->ne[1];           \
28    const uint32_t ne2 = dst->ne[2];           \
29    const uint32_t ne3 = dst->ne[3];           \
30                                               \
31    const uint32_t nb1 = dst->nb[1];           \
32    const uint32_t nb2 = dst->nb[2];           \
33    const uint32_t nb3 = dst->nb[3];           \
34                                               \
35    const uint32_t nr = ne1 * ne2 * ne3;
36
37struct htp_fill_context {
38    struct htp_ops_context * octx;
39    uint32_t nrows_per_thread;
40    uint32_t total_rows;  // ne1 * ne2 * ne3
41    uint32_t row_start;
42    bool     opt_path;
43    HVX_Vector splat_vec;
44    uint32_t   elem_size;
45};
46
47static void fill_thread(unsigned int nth, unsigned int ith, void * data) {
48    const struct htp_fill_context * fctx = (const struct htp_fill_context *) data;
49    struct htp_ops_context        * octx = fctx->octx;
50    fill_preamble;
51
52    // Parallelise over the flat row index spanning ne1*ne2*ne3
53    const uint32_t ir0 = fctx->row_start + fctx->nrows_per_thread * ith;
54    const uint32_t ir1 = MIN(ir0 + fctx->nrows_per_thread, fctx->row_start + fctx->total_rows);
55
56    if (ir0 >= ir1) {
57        return;
58    }
59
60    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
61    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir0);
62
63    if (fctx->opt_path) {
64        // Opt path: tensor is fully contiguous, treat as flat array
65        const uint32_t elem_start = ir0 * ne0;
66        const uint32_t elem_end = ir1 * ne0;
67        uint8_t * dst_ptr = (uint8_t *) dst->data + elem_start * fctx->elem_size;
68        hvx_splat_u(dst_ptr, fctx->splat_vec, elem_end - elem_start, fctx->elem_size);
69    } else {
70        // Non-contiguous path: must respect strides
71        for (uint32_t ir = ir0; ir < ir1; ++ir) {
72            const uint32_t i1 = ir % ne1;
73            const uint32_t i2 = (ir / ne1) % ne2;
74            const uint32_t i3 = ir / (ne1 * ne2);
75            uint8_t * dst_ptr = (uint8_t *) dst->data + i1*nb1 + i2*nb2 + i3*nb3;
76            hvx_splat_u(dst_ptr, fctx->splat_vec, ne0, fctx->elem_size);
77        }
78    }
79
80    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir1);
81    FARF(HIGH, "fill %u/%u: rows %u:%u\n", ith, nth, ir0, ir1);
82}
83
84int op_fill(struct htp_ops_context * octx) {
85    fill_preamble;
86
87    if (dst->type != HTP_TYPE_F32 && dst->type != HTP_TYPE_F16) {
88        return HTP_STATUS_NO_SUPPORT;
89    }
90
91    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
92        return HTP_STATUS_OK;
93    }
94
95    uint32_t row_start = 0;
96    uint32_t nrows     = nr;
97
98    if (octx->ctx->mdev.count > 1) {
99        const uint32_t row_size = nb1;
100        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
101        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(nr, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
102        row_start = range.start;
103        nrows     = range.count;
104    }
105
106    if (nrows == 0) {
107        return HTP_STATUS_OK;
108    }
109
110    // nr = ne1*ne2*ne3 (flat row count across all outer dims); parallelise over it.
111    const uint32_t n_threads = octx->n_threads;
112
113    // Optimize if fully contiguous: skip stride arithmetic, treat as flat array
114    const bool opt_path = (nb2 == nb1 * ne1) && (nb3 == nb2 * ne2);
115
116    FARF(HIGH, "fill: (%ux%ux%ux%u) type=%u opt=%d\n",
117         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], dst->type, (int) opt_path);
118
119    float val_f32 = 0.f;
120    memcpy(&val_f32, &octx->op_params[0], sizeof(float));
121
122    struct htp_fill_context fctx = {
123        .octx             = octx,
124        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
125        .total_rows       = nrows,
126        .row_start        = row_start,
127        .opt_path         = opt_path,
128    };
129
130    switch (dst->type) {
131    case HTP_TYPE_F32:
132        fctx.splat_vec = hvx_vec_splat_f32(val_f32);
133        fctx.elem_size = sizeof(float);
134        break;
135    case HTP_TYPE_F16:
136        fctx.splat_vec = hvx_vec_splat_f16((_Float16) val_f32);
137        fctx.elem_size = sizeof(_Float16);
138        break;
139    default:
140        return HTP_STATUS_NO_SUPPORT;
141    }
142
143    work_queue_run(octx->ctx->work_queue, fill_thread, &fctx, n_threads);
144
145    return HTP_STATUS_OK;
146}