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.5 KiB138 linesraw
1#ifndef HTP_TENSOR_H
2#define HTP_TENSOR_H
3
4#include <stdint.h>
5#include <stdbool.h>
6#include "htp-ops.h"
7#include "hex-bitmap.h"
8#include "hex-common.h"
9#include "hex-fastdiv.h"
10
11enum {
12    HTP_TENSOR_MDEV_LINE_SIZE = 128,
13};
14
15struct htp_tensor_mdev_range {
16    uint32_t start;
17    uint32_t count;
18};
19
20static inline void * htp_tensor_data(const struct htp_tensor * t) {
21    return (void *) (uintptr_t) t->data;
22}
23
24static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) {
25    return (uint32_t *) &t->flags;
26}
27
28static inline bool htp_tensor_is_contiguous(const struct htp_tensor * t, uint32_t type_size) {
29    uint32_t next_nb = type_size;
30    if (t->ne[0] != 1 && t->nb[0] != next_nb) {
31        return false;
32    }
33    next_nb *= t->ne[0];
34    for (int i = 1; i < HTP_OP_MAX_DIMS; i++) {
35        if (t->ne[i] != 1 && t->nb[i] != next_nb) {
36            return false;
37        }
38        next_nb *= t->ne[i];
39    }
40    return true;
41}
42
43static inline bool htp_tensor_is_permuted(const struct htp_tensor * t) {
44    return t->nb[0] > t->nb[1] || t->nb[1] > t->nb[2] || t->nb[2] > t->nb[3];
45}
46
47static inline bool htp_tensor_mdev_data_aligned(const struct htp_tensor * t) {
48    return ((uintptr_t) t->data & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0;
49}
50
51static inline bool htp_tensor_can_row_partition(const struct htp_tensor * t, uint32_t elem_size) {
52    if (!htp_tensor_mdev_data_aligned(t)) {
53        return false;
54    }
55    if (t->ne[0] != 1 && t->nb[0] != elem_size) {
56        return false;
57    }
58    if (htp_tensor_is_permuted(t)) {
59        return false;
60    }
61    if (t->ne[1] > 1 && (t->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
62    if (t->ne[2] > 1 && (t->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
63    if (t->ne[3] > 1 && (t->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) != 0) return false;
64    return true;
65}
66
67static inline bool htp_tensor_mdev_rows_per_chunk(const struct htp_tensor * t, uint32_t elem_size, uint32_t row_size, uint32_t * rows_per_chunk) {
68    *rows_per_chunk = 0;
69
70    if (!htp_tensor_mdev_data_aligned(t)) {
71        return false;
72    }
73    if (t->ne[0] != 1 && t->nb[0] != elem_size) {
74        return false;
75    }
76    if (htp_tensor_is_permuted(t)) {
77        return false;
78    }
79    if (t->ne[1] > 1 && (t->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0 &&
80        (t->ne[2] <= 1 || (t->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0) &&
81        (t->ne[3] <= 1 || (t->nb[3] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0)) {
82        *rows_per_chunk = 1;
83        return true;
84    }
85    if (t->nb[1] == row_size &&
86        (t->ne[2] <= 1 || t->nb[2] == t->nb[1] * t->ne[1]) &&
87        (t->ne[3] <= 1 || t->nb[3] == t->nb[2] * t->ne[2])) {
88        *rows_per_chunk = (row_size > 0) ? (HTP_TENSOR_MDEV_LINE_SIZE / hex_gcd_u32(row_size, HTP_TENSOR_MDEV_LINE_SIZE)) : 1;
89        return true;
90    }
91    return false;
92}
93
94static inline struct htp_tensor_mdev_range htp_tensor_mdev_partition(uint32_t total_units, uint32_t units_per_chunk, uint32_t mdev_idx, uint32_t mdev_count, const struct fastdiv_values * mdev_count_div) {
95    struct htp_tensor_mdev_range range = { 0, total_units };
96
97    if (mdev_count <= 1) {
98        return range;
99    }
100
101    if (units_per_chunk == 0) {
102        range.start = (mdev_idx == 0) ? 0 : total_units;
103        range.count = (mdev_idx == 0) ? total_units : 0;
104        return range;
105    }
106
107    const uint32_t total_chunks = total_units / units_per_chunk;
108    if (total_chunks < mdev_count) {
109        range.start = (mdev_idx == 0) ? 0 : total_units;
110        range.count = (mdev_idx == 0) ? total_units : 0;
111        return range;
112    }
113
114    const uint32_t chunks_per_mdev = fastdiv(total_chunks + mdev_count - 1, mdev_count_div);
115    range.start = MIN(mdev_idx * chunks_per_mdev * units_per_chunk, total_units);
116    if (mdev_idx == mdev_count - 1) {
117        range.count = total_units - range.start;
118    } else {
119        range.count = MIN(chunks_per_mdev * units_per_chunk, total_units - range.start);
120    }
121    return range;
122}
123
124static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) {
125    switch (type) {
126        case HTP_TYPE_F32:  return ne00 * 4;
127        case HTP_TYPE_F16:  return ne00 * 2;
128        case HTP_TYPE_Q8_0: return (ne00 / 32) * 34;
129        default:            return 0;
130    }
131}
132
133struct htp_context;
134void htp_flush_dirty_ranges(struct htp_context * ctx);
135void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
136void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n);
137
138#endif // HTP_TENSOR_H