hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

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)eafe15a5e2026-09-11

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#include <stdatomic.h>
8#include <math.h>
9#include <string.h>
10
11#define GGML_COMMON_DECL_C
12#include "ggml-common.h"
13#include "htp-ctx.h"
14#include "htp-ops.h"
15#include "hvx-utils.h"
16#include "htp-tensor.h"
17#include "hex-dma.h"
18#include "hex-profile.h"
19#include "allreduce-ops.h"
20#include "htp-fence.h"
21
22struct htp_allreduce_context {
23    struct htp_ops_context * octx;
24    uint32_t n_ranks;
25    uint32_t n_dsts;
26    uint32_t nelem;
27    uint32_t ne0;
28    uint32_t ne1;
29    uint32_t row_size_aligned;
30    uint32_t rank_elem_start;
31    uint32_t rank_nelem;
32    uint32_t elems_per_thread;
33    uint32_t block_elems;
34    uint32_t vtcm_size_per_thread;
35    bool     is_row_bcast;
36    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];
37    uint8_t * dst_spad_base;
38    uint8_t * res_spad_base;
39};
40
41#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD)                                             \
42static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                       \
43    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                      \
44    struct htp_ops_context * octx = actx->octx;                                                                       \
45                                                                                                                      \
46    const uint32_t n_ranks     = actx->n_ranks;                                                                       \
47    const uint32_t n_dsts      = actx->n_dsts;                                                                        \
48    const uint32_t block_elems = actx->block_elems;                                                                   \
49                                                                                                                      \
50    const uint32_t dr  = actx->elems_per_thread;                                                                      \
51    const uint32_t ir0 = actx->rank_elem_start + dr * ith;                                                            \
52    const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem);                                     \
53    if (ir0 >= ir1) return;                                                                                           \
54                                                                                                                      \
55    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                            \
56    dma_queue * q = octx->ctx->dma[ith];                                                                              \
57                                                                                                                      \
58    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                 \
59    for (uint32_t s = 0; s < n_ranks; s++) {                                                                          \
60        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                               \
61    }                                                                                                                 \
62    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                               \
63    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL;            \
64                                                                                                                      \
65    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                          \
66    uint32_t ir_prefetch = ir0;                                                                                       \
67    int spad_idx = 0;                                                                                                 \
68                                                                                                                      \
69    for (int k = 0; k < 2 && ir_prefetch < ir1; k++) {                                                                \
70        uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch);                                                     \
71        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                                \
72        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                      \
73        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
74            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir_prefetch * sizeof(TYPE);                           \
75            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0);                       \
76        }                                                                                                             \
77        for (uint32_t s = 0; s < n_ranks; s++) {                                                                      \
78            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                               \
79            const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE);                \
80            dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1);                       \
81        }                                                                                                             \
82        if (HAS_ADD) {                                                                                                \
83            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                                  \
84            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);      \
85            dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1);                       \
86        }                                                                                                             \
87        ir_prefetch += cur_elems;                                                                                     \
88        spad_idx ^= 1;                                                                                                \
89    }                                                                                                                 \
90                                                                                                                      \
91    for (uint32_t ir = ir0; ir < ir1; ) {                                                                             \
92        uint32_t cur_elems = MIN(block_elems, ir1 - ir);                                                              \
93        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                                \
94        uint8_t * d_spad = NULL;                                                                                      \
95        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
96            d_spad = (uint8_t *) dma_queue_pop(q).src;                                                                \
97        }                                                                                                             \
98        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                    \
99        for (uint32_t s = 0; s < n_ranks; s++) {                                                                      \
100            s_spad[s] = (uint8_t *) dma_queue_pop(q).dst;                                                             \
101        }                                                                                                             \
102        uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(q).dst : NULL;                                         \
103        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                             \
104        HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems);                                                          \
105        for (uint32_t s = 2; s < n_ranks; s++) {                                                                      \
106            HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems);                                                         \
107        }                                                                                                             \
108        if (HAS_ADD) {                                                                                                \
109            HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems);                                                            \
110        }                                                                                                             \
111        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                              \
112        for (uint32_t d = 0; d < n_dsts; d++) {                                                                       \
113            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir * sizeof(TYPE);                                    \
114            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1);                       \
115        }                                                                                                             \
116        if (ir_prefetch < ir1) {                                                                                      \
117            uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch);                                                \
118            size_t   next_bytes = next_elems * sizeof(TYPE);                                                          \
119            for (uint32_t s = 0; s < n_ranks; s++) {                                                                  \
120                const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE);           \
121                dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1);            \
122            }                                                                                                         \
123            if (HAS_ADD) {                                                                                            \
124                const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \
125                dma_queue_push(q, dma_make_ptr(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1);               \
126            }                                                                                                         \
127            ir_prefetch += next_elems;                                                                                \
128        }                                                                                                             \
129        ir += cur_elems;                                                                                              \
130    }                                                                                                                 \
131    dma_queue_flush(q);                                                                                               \
132}
133
134DEFINE_ALLREDUCE_THREAD_DMA_1D(f16,     __fp16, hvx_add_f16_aaa, 0)
135DEFINE_ALLREDUCE_THREAD_DMA_1D(f32,     float,  hvx_add_f32_aaa, 0)
136DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f16, __fp16, hvx_add_f16_aaa, 1)
137DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f32, float,  hvx_add_f32_aaa, 1)
138
139#define DEFINE_ALLREDUCE_THREAD_DMA_2D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD, IS_ROW_BCAST)                                                           \
140static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                                                   \
141    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                                                  \
142    struct htp_ops_context * octx = actx->octx;                                                                                                   \
143                                                                                                                                                  \
144    const uint32_t n_ranks          = actx->n_ranks;                                                                                              \
145    const uint32_t n_dsts           = actx->n_dsts;                                                                                               \
146    const uint32_t ne0              = actx->ne0;                                                                                                  \
147    const uint32_t block_rows       = actx->block_elems;                                                                                          \
148    const uint32_t row_size_aligned = actx->row_size_aligned;                                                                                     \
149    const uint32_t row_bytes        = ne0 * sizeof(TYPE);                                                                                         \
150                                                                                                                                                  \
151    const uint32_t dr  = actx->elems_per_thread;                                                                                                  \
152    const uint32_t r0  = actx->rank_elem_start + dr * ith;                                                                                        \
153    const uint32_t r1  = MIN(r0 + dr, actx->rank_elem_start + actx->rank_nelem);                                                                  \
154    if (r0 >= r1) return;                                                                                                                         \
155                                                                                                                                                  \
156    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                                        \
157    dma_queue * q = octx->ctx->dma[ith];                                                                                                          \
158                                                                                                                                                  \
159    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                                             \
160    for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                      \
161        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                                                           \
162    }                                                                                                                                             \
163    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                                                           \
164    uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + (ith * actx->vtcm_size_per_thread))) : NULL; \
165                                                                                                                                                  \
166    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                                                      \
167    uint32_t r_prefetch = r0;                                                                                                                     \
168    int spad_idx = 0;                                                                                                                             \
169                                                                                                                                                  \
170    for (int k = 0; k < 2 && r_prefetch < r1; k++) {                                                                                              \
171        uint32_t cur_rows = MIN(block_rows, r1 - r_prefetch);                                                                                     \
172        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                                                  \
173        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
174            uint8_t * d_ddr  = (uint8_t *) octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1];                                               \
175            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0);                                 \
176        }                                                                                                                                         \
177        for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
178            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                                                           \
179            const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                      \
180            dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows);                           \
181        }                                                                                                                                         \
182        if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                           \
183            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                                                              \
184            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];                  \
185            dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows);                 \
186        }                                                                                                                                         \
187        r_prefetch += cur_rows;                                                                                                                   \
188        spad_idx ^= 1;                                                                                                                            \
189    }                                                                                                                                             \
190                                                                                                                                                  \
191    for (uint32_t r = r0; r < r1; ) {                                                                                                             \
192        uint32_t cur_rows = MIN(block_rows, r1 - r);                                                                                              \
193        uint8_t * d_spad = NULL;                                                                                                                  \
194        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
195            d_spad = (uint8_t *) dma_queue_pop(q).src;                                                                                            \
196        }                                                                                                                                         \
197        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                                                \
198        for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
199            s_spad[s] = (uint8_t *) dma_queue_pop(q).dst;                                                                                         \
200        }                                                                                                                                         \
201        uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(q).dst : NULL;                                                  \
202        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                          \
203        for (uint32_t row = 0; row < cur_rows; row++) {                                                                                           \
204            uint8_t * d_row = d_spad + row * row_size_aligned;                                                                                    \
205            const uint8_t * s0_row = s_spad[0] + row * row_size_aligned;                                                                          \
206            const uint8_t * s1_row = s_spad[1] + row * row_size_aligned;                                                                          \
207            HVX_ADD_FN(d_row, s0_row, s1_row, ne0);                                                                                               \
208            for (uint32_t s = 2; s < n_ranks; s++) {                                                                                              \
209                const uint8_t * ss_row = s_spad[s] + row * row_size_aligned;                                                                      \
210                HVX_ADD_FN(d_row, d_row, ss_row, ne0);                                                                                            \
211            }                                                                                                                                     \
212            if (HAS_ADD) {                                                                                                                        \
213                const uint8_t * res_row = IS_ROW_BCAST ? res_spad_base : (r_spad + row * row_size_aligned);                                       \
214                HVX_ADD_FN(d_row, d_row, res_row, ne0);                                                                                           \
215            }                                                                                                                                     \
216        }                                                                                                                                         \
217        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                           \
218        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
219            uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r * octx->dsts[d]->nb[1];                                                         \
220            dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows);                          \
221        }                                                                                                                                         \
222        if (r_prefetch < r1) {                                                                                                                    \
223            uint32_t next_rows = MIN(block_rows, r1 - r_prefetch);                                                                                \
224            for (uint32_t s = 0; s < n_ranks; s++) {                                                                                              \
225                const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                 \
226                dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows);                  \
227            }                                                                                                                                     \
228            if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                       \
229                const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];             \
230                dma_queue_push(q, dma_make_ptr(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows);           \
231            }                                                                                                                                     \
232            r_prefetch += next_rows;                                                                                                              \
233        }                                                                                                                                         \
234        r += cur_rows;                                                                                                                            \
235    }                                                                                                                                             \
236    dma_queue_flush(q);                                                                                                                           \
237}
238
239DEFINE_ALLREDUCE_THREAD_DMA_2D(f16,           __fp16, hvx_add_f16_aaa, 0, 0)
240DEFINE_ALLREDUCE_THREAD_DMA_2D(f32,           float,  hvx_add_f32_aaa, 0, 0)
241DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f16,       __fp16, hvx_add_f16_aaa, 1, 0)
242DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f32,       float,  hvx_add_f32_aaa, 1, 0)
243DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f16, __fp16, hvx_add_f16_aaa, 1, 1)
244DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f32, float,  hvx_add_f32_aaa, 1, 1)
245
246static int validate_allreduce(
247    struct htp_ops_context * octx,
248    const struct htp_allreduce_kernel_params * kparams,
249    uint32_t n_ranks
250) {
251    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
252        return HTP_STATUS_INVAL_PARAMS;
253    }
254
255    if (kparams->vtcm_size_per_thread <= 0 || kparams->vtcm_size <= 0) {
256        return HTP_STATUS_INVAL_PARAMS;
257    }
258
259    const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
260    const size_t n_vtcm_buffers = htp_allreduce_vtcm_buffer_count(
261        n_ranks, octx->n_threads, has_add, kparams->is_row_bcast != 0);
262    const size_t vtcm_size = n_vtcm_buffers * (size_t) kparams->vtcm_size_per_thread;
263    if (vtcm_size != (size_t) kparams->vtcm_size) {
264        return HTP_STATUS_INVAL_PARAMS;
265    }
266    if (vtcm_size > octx->ctx->vtcm_size) {
267        return HTP_STATUS_VTCM_TOO_SMALL;
268    }
269
270    if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
271        return HTP_STATUS_NO_SUPPORT;
272    }
273
274    return HTP_STATUS_OK;
275}
276
277int op_allreduce(struct htp_ops_context * octx) {
278    if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
279        return HTP_STATUS_OK;
280    }
281
282    const struct htp_allreduce_kernel_params * kparams = (const struct htp_allreduce_kernel_params *) octx->kernel_params;
283    const struct htp_tensor * dst = octx->dst;
284
285    const uint32_t rank    = (uint32_t) kparams->rank;
286    const uint32_t n_ranks = (uint32_t) kparams->n_ranks;
287
288    if (n_ranks < 2 || n_ranks > HTP_ALLREDUCE_MAX_RANKS || rank >= n_ranks) {
289        return HTP_STATUS_INVAL_PARAMS;
290    }
291
292    const uint32_t fence_seq_entry = (uint32_t) octx->op_params[0];
293    const uint32_t fence_seq_exit  = (uint32_t) octx->op_params[1];
294
295    const struct htp_tensor * my_sync = octx->src[n_ranks + rank];
296    atomic_uint * my_fence = (atomic_uint *) (uintptr_t) my_sync->data;
297
298    const int status = validate_allreduce(octx, kparams, n_ranks);
299    if (status != HTP_STATUS_OK) {
300        if (status == HTP_STATUS_NO_SUPPORT) {
301            FARF(ERROR, "ggml-hex: allreduce unsupported type %d : rank %u\n", dst->type, rank);
302        }
303        htp_fence_write(my_fence, fence_seq_exit, status);
304        return status;
305    }
306
307    const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
308    const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
309
310    // 1. Entry Barrier: Synchronize all ranks before reading
311    struct htp_thread_trace * tr0 = &octx->ctx->trace[0];
312    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
313
314    htp_fence_write(my_fence, fence_seq_entry, octx->status);
315
316    for (uint32_t j = 0; j < n_ranks; j++) {
317        if (j == rank) continue;
318        const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
319        atomic_uint * peer_fence = (atomic_uint *) (uintptr_t) peer_sync->data;
320        uint64_t spins = 0;
321        while (1) {
322            uint32_t peer_seq;
323            uint32_t peer_status;
324            htp_fence_read(peer_fence, &peer_seq, &peer_status);
325            if ((int32_t)(peer_seq - fence_seq_entry) >= 0) {
326                if (peer_status > HTP_STATUS_OK) {
327                    FARF(ERROR, "ggml-hex: allreduce entry peer %u failed with status %u\n", j, peer_status);
328                    htp_fence_write(my_fence, fence_seq_exit, peer_status);
329                    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
330                    return peer_status;
331                }
332                break;
333            }
334            if (++spins > HTP_FENCE_TIMEOUT) {
335                FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT : rank %u waiting on %u fence %p seq 0x%x peer-seq 0x%x\n",
336                     rank, j, peer_fence, fence_seq_entry, peer_seq);
337                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_INTERNAL_ERR);
338                htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
339                return HTP_STATUS_INTERNAL_ERR;
340            }
341            hex_pause();
342        }
343    }
344    asm volatile ("syncht" : : : "memory");
345
346    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
347
348    // 2. Multi-threaded Reduction across assigned rank chunk
349    if (nelem > 0) {
350        const uint32_t n_threads            = (uint32_t) kparams->n_threads;
351        const uint32_t block_elems          = (uint32_t) kparams->block_elems;
352        const uint32_t elems_per_thread     = (uint32_t) kparams->elems_per_thread;
353        const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread;
354
355        struct htp_allreduce_context actx;
356        actx.octx                 = octx;
357        actx.n_ranks              = n_ranks;
358        actx.n_dsts               = (uint32_t) kparams->n_dsts ? (uint32_t) kparams->n_dsts : n_ranks;
359        actx.nelem                = nelem;
360        actx.ne0                  = (uint32_t) kparams->ne0;
361        actx.ne1                  = (uint32_t) kparams->ne1;
362        actx.row_size_aligned     = (uint32_t) kparams->row_size_aligned;
363        actx.rank_elem_start      = (uint32_t) kparams->rank_elem_start;
364        actx.rank_nelem           = (uint32_t) kparams->rank_nelem;
365        actx.elems_per_thread     = elems_per_thread;
366        actx.block_elems          = block_elems;
367        actx.vtcm_size_per_thread = vtcm_size_per_thread;
368        actx.is_row_bcast         = (kparams->is_row_bcast != 0);
369
370        work_queue_func_t reduce_fun = NULL;
371        switch (kparams->kernel_type) {
372            case HTP_ALLREDUCE_KERNEL_DMA_1D:
373                if (has_add) {
374                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_add_f16 : allreduce_thread_dma_1d_add_f32;
375                } else {
376                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_f16 : allreduce_thread_dma_1d_f32;
377                }
378                break;
379            case HTP_ALLREDUCE_KERNEL_DMA_2D:
380                if (has_add) {
381                    if (kparams->is_row_bcast) {
382                        reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_bcast_f16 : allreduce_thread_dma_2d_add_bcast_f32;
383                    } else {
384                        reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_f16 : allreduce_thread_dma_2d_add_f32;
385                    }
386                } else {
387                    reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_f16 : allreduce_thread_dma_2d_f32;
388                }
389                break;
390            default:
391                FARF(ERROR, "ggml-hex: allreduce unsupported kernel %d : rank %u\n", kparams->kernel_type, rank);
392                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_NO_SUPPORT);
393                return HTP_STATUS_NO_SUPPORT;
394        }
395
396        uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
397        for (uint32_t s = 0; s < n_ranks; s++) {
398            actx.src_spad_base[s] = vtcm_ptr;
399            vtcm_ptr += n_threads * vtcm_size_per_thread;
400        }
401        actx.dst_spad_base = vtcm_ptr;
402        vtcm_ptr += n_threads * vtcm_size_per_thread;
403        if (has_add) {
404            actx.res_spad_base = vtcm_ptr;
405            vtcm_ptr += (actx.is_row_bcast ? 1 : n_threads) * vtcm_size_per_thread;
406        }
407
408        if (has_add && actx.is_row_bcast) {
409            const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data;
410            const uint32_t row_bytes = actx.ne0 * (dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float));
411            dma_queue * q = octx->ctx->dma[0];
412            dma_queue_push(q, dma_make_ptr(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1);
413            dma_queue_pop(q);
414        }
415
416        work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads);
417    }
418
419    // 4. Exit Barrier: Synchronize all ranks after writing
420    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
421
422    htp_fence_write(my_fence, fence_seq_exit, octx->status);
423
424    for (uint32_t j = 0; j < n_ranks; j++) {
425        if (j == rank) continue;
426        const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
427        atomic_uint * peer_fence = (atomic_uint *) (uintptr_t) peer_sync->data;
428        uint64_t spins = 0;
429        while (1) {
430            uint32_t peer_seq;
431            uint32_t peer_status;
432            htp_fence_read(peer_fence, &peer_seq, &peer_status);
433            if ((int32_t)(peer_seq - fence_seq_exit) >= 0) {
434                if (peer_status > HTP_STATUS_OK) {
435                    FARF(ERROR, "ggml-hex: allreduce exit peer %u failed with status %u\n", j, peer_status);
436                    htp_fence_write(my_fence, fence_seq_exit, peer_status);
437                    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
438                    return peer_status;
439                }
440                break;
441            }
442            if (++spins > HTP_FENCE_TIMEOUT) {
443                FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT : rank %u waiting on %u fence %p seq 0x%x peer-seq 0x%x\n",
444                     rank, j, peer_fence, fence_seq_exit, peer_seq);
445                htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_INTERNAL_ERR);
446                htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
447                return HTP_STATUS_INTERNAL_ERR;
448            }
449            hex_pause();
450        }
451    }
452    asm volatile ("syncht" : : : "memory");
453
454    htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
455
456    return octx->status;
457}