char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Neo Zhangsycl : Support DSv4 OPs: LIGHTNING_INDEXER,DSV4_HC_COMB,DSV4_HC_POST,DSV4_HC_PRE (#26568)31558dbb7

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
8.6 KiB197 linesraw
1#include "lightning-indexer.hpp"
2#include "dequantize.hpp"
3
4static void lightning_indexer_f32_sycl(
5        const char * q, const char * k, const char * w, const char * m, float * dst,
6        int64_t n_embd, int64_t n_head, int64_t n_batch, int64_t n_stream, int64_t n_kv,
7        int64_t nem3,
8        int64_t nbq1, int64_t nbq2, int64_t nbq3,
9        int64_t nbk2, int64_t nbk3,
10        int64_t nbw1, int64_t nbw3,
11        int64_t nbm1, int64_t nbm3,
12        int64_t nb1, int64_t nb3,
13        ggml_type k_type,
14        queue_ptr stream) {
15
16    constexpr int64_t LANES = WARP_SIZE;
17    constexpr int64_t ELEMS_PER_LANE = 8;
18    constexpr int64_t ROWS_PER_BLOCK = 4;
19    constexpr int64_t BLOCK_SIZE = ROWS_PER_BLOCK * LANES;
20
21    const int64_t n_rows = n_batch * n_stream * n_kv;
22    const int64_t n_blocks = (n_rows + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK;
23
24    stream->parallel_for(
25        sycl::nd_range<1>(
26            sycl::range<1>(n_blocks * BLOCK_SIZE),
27            sycl::range<1>(BLOCK_SIZE)),
28        [=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
29            const int64_t ir   = item.get_global_id(0);
30            const int64_t lane = ir % LANES;
31            const int64_t row  = ir / LANES;
32            if (row >= n_rows) {
33                return;
34            }
35
36            const int64_t i_bs     = row / n_kv;
37            const int64_t i_kv     = row % n_kv;
38            const int64_t i_batch  = i_bs / n_stream;
39            const int64_t i_stream = i_bs % n_stream;
40
41            // load K row slice into registers (row is contiguous, nbk0 == type size)
42            const char * k_base = k + i_kv*nbk2 + i_stream*nbk3;
43            float k_local[ELEMS_PER_LANE];
44            if (k_type == GGML_TYPE_F16) {
45                const sycl::half * k_row = (const sycl::half *) k_base;
46#pragma unroll
47                for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
48                    k_local[j] = static_cast<float>(k_row[lane*ELEMS_PER_LANE + j]);
49                }
50            } else if (k_type == GGML_TYPE_F32) {
51                const float * k_row = (const float *) k_base;
52#pragma unroll
53                for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
54                    k_local[j] = k_row[lane*ELEMS_PER_LANE + j];
55                }
56            } else {
57                const int64_t lane_base = lane * ELEMS_PER_LANE;
58                switch (k_type) {
59                    case GGML_TYPE_BF16: {
60                        const sycl::ext::oneapi::bfloat16 * k_row = (const sycl::ext::oneapi::bfloat16 *) k_base;
61#pragma unroll
62                        for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
63                            k_local[j] = static_cast<float>(k_row[lane_base + j]);
64                        }
65                    } break;
66                    case GGML_TYPE_Q4_0:
67                    case GGML_TYPE_Q4_1:
68                    case GGML_TYPE_Q5_0:
69                    case GGML_TYPE_Q5_1: {
70#pragma unroll
71                        for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
72                            const int64_t idx = lane_base + j;
73                            const int64_t ib  = idx / QK4_0;
74                            const int iqs     = idx % (QK4_0/2);
75                            dfloat2 kv;
76                            if (k_type == GGML_TYPE_Q4_0) {
77                                dequantize_q4_0(k_base, ib, iqs, kv);
78                            } else if (k_type == GGML_TYPE_Q4_1) {
79                                dequantize_q4_1(k_base, ib, iqs, kv);
80                            } else if (k_type == GGML_TYPE_Q5_0) {
81                                dequantize_q5_0(k_base, ib, iqs, kv);
82                            } else {
83                                dequantize_q5_1(k_base, ib, iqs, kv);
84                            }
85                            k_local[j] = (idx % QK4_0) < (QK4_0/2) ? static_cast<float>(kv.x()) : static_cast<float>(kv.y());
86                        }
87                    } break;
88                    case GGML_TYPE_Q8_0: {
89#pragma unroll
90                        for (int64_t pair = 0; pair < ELEMS_PER_LANE / 2; ++pair) {
91                            const int64_t elem0 = lane_base + 2 * pair;
92                            dfloat2 kv;
93                            dequantize_q8_0(k_base, elem0 / QK8_0, elem0 % QK8_0, kv);
94                            k_local[2 * pair + 0] = static_cast<float>(kv.x());
95                            k_local[2 * pair + 1] = static_cast<float>(kv.y());
96                        }
97                    } break;
98                    case GGML_TYPE_IQ4_NL: {
99#pragma unroll
100                        for (int64_t pair = 0; pair < ELEMS_PER_LANE / 2; ++pair) {
101                            const int64_t elem0 = lane_base + 2 * pair;
102                            dfloat2 kv;
103                            dequantize_iq4_nl(k_base, elem0 / QK4_NL, elem0 % QK4_NL, kv);
104                            k_local[2 * pair + 0] = static_cast<float>(kv.x());
105                            k_local[2 * pair + 1] = static_cast<float>(kv.y());
106                        }
107                    } break;
108                    default:
109#pragma unroll
110                        for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
111                            k_local[j] = 0.0f;
112                        }
113                        break;
114                }
115            }
116
117            const char  * q_base = q + i_batch*nbq2 + i_stream*nbq3;
118            const float * w_base = (const float *) (w + i_batch*nbw1 + i_stream*nbw3);
119
120            float score = 0.0f;
121            for (int64_t h = 0; h < n_head; ++h) {
122                const float * q_row = (const float *) (q_base + h*nbq1);
123                float dot = 0.0f;
124#pragma unroll
125                for (int64_t j = 0; j < ELEMS_PER_LANE; ++j) {
126                    const int64_t i = lane*ELEMS_PER_LANE + j;
127                    if (i < n_embd) {
128                        dot += q_row[i] * k_local[j];
129                    }
130                }
131                dot = sycl::reduce_over_group(item.get_sub_group(), dot, sycl::plus<float>());
132                if (lane == 0) {
133                    score += sycl::max(dot, 0.0f) * w_base[h];
134                }
135            }
136
137            if (lane == 0) {
138                const sycl::half * m_base = (const sycl::half *) (m + i_batch*nbm1 + (i_stream % nem3)*nbm3);
139                // flat-index store: storing through a strided base pointer
140                // hangs/misroutes writes on this stack when n_batch*n_stream > 1
141                const int64_t dst_idx = i_kv + i_batch*(nb1/sizeof(float)) + i_stream*(nb3/sizeof(float));
142                dst[dst_idx] = score + static_cast<float>(m_base[i_kv]);
143            }
144        });
145}
146
147void ggml_sycl_op_lightning_indexer(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
148    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/4);
149    const ggml_tensor * q = dst->src[0];
150    const ggml_tensor * k = dst->src[1];
151    const ggml_tensor * w = dst->src[2]; // weights
152    const ggml_tensor * m = dst->src[3]; // mask
153
154    GGML_ASSERT(dst->type == GGML_TYPE_F32);
155    GGML_ASSERT(  q->type == GGML_TYPE_F32);
156    GGML_ASSERT(  w->type == GGML_TYPE_F32);
157    GGML_ASSERT(  m->type == GGML_TYPE_F16);
158    GGML_ASSERT(k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_BF16 ||
159                k->type == GGML_TYPE_Q8_0 || k->type == GGML_TYPE_Q5_1 || k->type == GGML_TYPE_Q5_0 ||
160                k->type == GGML_TYPE_Q4_1 || k->type == GGML_TYPE_Q4_0 || k->type == GGML_TYPE_IQ4_NL);
161
162    GGML_TENSOR_LOCALS(int64_t, neq, q, ne);
163    GGML_TENSOR_LOCALS(size_t,  nbq, q, nb);
164    GGML_TENSOR_LOCALS(int64_t, nek, k, ne);
165    GGML_TENSOR_LOCALS(size_t,  nbk, k, nb);
166    GGML_TENSOR_LOCALS(size_t,  nbw, w, nb);
167    GGML_TENSOR_LOCALS(int64_t, nem, m, ne);
168    GGML_TENSOR_LOCALS(size_t,  nbm, m, nb);
169    GGML_TENSOR_LOCALS(int64_t, ne, dst, ne);
170    GGML_TENSOR_LOCALS(size_t,  nb, dst, nb);
171
172    // input rows must be contiguous
173    GGML_ASSERT(nbq0 == ggml_type_size(q->type));
174    GGML_ASSERT(nbk0 == ggml_type_size(k->type));
175    GGML_ASSERT(nbm0 == ggml_type_size(m->type));
176    GGML_ASSERT(nb0  == ggml_type_size(dst->type));
177
178    const int64_t n_embd   = neq0;
179    const int64_t n_head   = neq1;
180    const int64_t n_batch  = neq2;
181    const int64_t n_stream = neq3;
182    const int64_t n_kv     = nek2;
183
184    GGML_ASSERT(n_embd == WARP_SIZE * 8);
185
186    lightning_indexer_f32_sycl(
187            (const char *) q->data, (const char *) k->data,
188            (const char *) w->data, (const char *) m->data, (float *) dst->data,
189            n_embd, n_head, n_batch, n_stream, n_kv, nem3,
190            nbq1, nbq2, nbq3,
191            nbk2, nbk3,
192            nbw1, nbw3,
193            nbm1, nbm3,
194            nb1, nb3,
195            k->type,
196            ctx.stream());
197}