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 →
1.1 KiB51 linesraw
1#ifndef ALLREDUCE_OPS_H
2#define ALLREDUCE_OPS_H
3
4#include <stdint.h>
5#include <stddef.h>
6#include <stdbool.h>
7
8#define HTP_ALLREDUCE_MAX_RANKS 4
9
10#ifdef __cplusplus
11extern "C" {
12#endif
13
14enum htp_allreduce_kernel_type {
15    HTP_ALLREDUCE_KERNEL_UNSUPPORTED = 0,
16    HTP_ALLREDUCE_KERNEL_DMA_1D,
17    HTP_ALLREDUCE_KERNEL_DMA_2D,
18};
19
20static inline size_t htp_allreduce_vtcm_buffer_count(
21    uint32_t n_ranks,
22    uint32_t n_threads,
23    bool has_add,
24    bool is_row_bcast
25) {
26    return (size_t) (n_ranks + 1) * n_threads + (has_add ? (is_row_bcast ? 1 : n_threads) : 0);
27}
28
29struct htp_allreduce_kernel_params {
30    int32_t rank;
31    int32_t n_ranks;
32    int32_t n_threads;
33    int32_t block_elems;          // 1D: block_elems, 2D: block_rows
34    int32_t elems_per_thread;     // 1D: nelem_per_thread, 2D: nrows_per_thread
35    int32_t vtcm_size_per_thread;
36    int32_t vtcm_size;
37    int32_t kernel_type;
38    int32_t ne0;
39    int32_t ne1;
40    int32_t row_size_aligned;
41    int32_t rank_elem_start;
42    int32_t rank_nelem;
43    int32_t n_dsts;
44    int32_t is_row_bcast;
45};
46
47#ifdef __cplusplus
48}
49#endif
50
51#endif /* ALLREDUCE_OPS_H */