char-fork/llama.cpp

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

Aman GuptaDeepseekV4: Add fused hyper-connection ops (#25585)0dc74e332

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.5 KiB294 linesraw
1#include "common.cuh"
2#include "dsv4-hc.cuh"
3
4
5static constexpr int DSV4_HC = 4;
6
7
8static __device__ void dsv4_hc_comb_norm_cols(float * comb, float eps) {
9    for (int idst = 0; idst < DSV4_HC; ++idst) {
10        float sum = eps;
11        for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
12            sum += comb[idst + DSV4_HC*isrc];
13        }
14
15        const float inv_sum = 1.0f / sum;
16        for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
17            comb[idst + DSV4_HC*isrc] *= inv_sum;
18        }
19    }
20}
21
22static __device__ void dsv4_hc_comb_norm_rows(float * comb, float eps) {
23    for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
24        float sum = eps;
25        for (int idst = 0; idst < DSV4_HC; ++idst) {
26            sum += comb[idst + DSV4_HC*isrc];
27        }
28
29        const float inv_sum = 1.0f / sum;
30        for (int idst = 0; idst < DSV4_HC; ++idst) {
31            comb[idst + DSV4_HC*isrc] *= inv_sum;
32        }
33    }
34}
35
36static __global__ void dsv4_hc_comb_f32(
37        const float * mixes,
38        const float * scale,
39        const float * base,
40        float * dst,
41        int64_t n_tokens,
42        int64_t sm0,
43        int64_t sm1,
44        int64_t ss0,
45        int64_t sb0,
46        int64_t sd0,
47        int64_t sd1,
48        int64_t sd2,
49        float eps,
50        int32_t n_iter) {
51    constexpr int comb_offset = 2*DSV4_HC;
52
53    ggml_cuda_pdl_lc();
54    const int64_t it = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
55
56    if (it >= n_tokens) {
57        return;
58    }
59
60    ggml_cuda_pdl_sync();
61
62    const float scale_comb = scale[2*ss0];
63    float comb[DSV4_HC*DSV4_HC];
64
65    for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
66        float max = -INFINITY;
67        for (int idst = 0; idst < DSV4_HC; ++idst) {
68            const int idx = idst + DSV4_HC*isrc;
69            const float v = mixes[(comb_offset + idx)*sm0 + it*sm1] * scale_comb + base[(comb_offset + idx)*sb0];
70            comb[idx] = v;
71            max = fmaxf(max, v);
72        }
73
74        float sum = 0.0f;
75        for (int idst = 0; idst < DSV4_HC; ++idst) {
76            const int idx = idst + DSV4_HC*isrc;
77            const float v = expf(comb[idx] - max);
78            comb[idx] = v;
79            sum += v;
80        }
81
82        const float inv_sum = 1.0f / sum;
83        for (int idst = 0; idst < DSV4_HC; ++idst) {
84            const int idx = idst + DSV4_HC*isrc;
85            comb[idx] = comb[idx] * inv_sum + eps;
86        }
87    }
88
89    dsv4_hc_comb_norm_cols(comb, eps);
90    for (int32_t i = 1; i < n_iter; ++i) {
91        dsv4_hc_comb_norm_rows(comb, eps);
92        dsv4_hc_comb_norm_cols(comb, eps);
93    }
94
95    for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
96        for (int idst = 0; idst < DSV4_HC; ++idst) {
97            const int idx = idst + DSV4_HC*isrc;
98            dst[idst*sd0 + isrc*sd1 + it*sd2] = comb[idx];
99        }
100    }
101}
102
103static __global__ void dsv4_hc_pre_f32(
104        const float * x,
105        const float * weights,
106        float * dst,
107        int64_t n_embd,
108        int64_t hc,
109        int64_t n_tokens,
110        int64_t sx0,
111        int64_t sx1,
112        int64_t sx2,
113        int64_t sw0,
114        int64_t sw1,
115        int64_t sd0,
116        int64_t sd1) {
117    ggml_cuda_pdl_lc();
118    const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
119    const int64_t nr = n_embd * n_tokens;
120
121    if (ir >= nr) {
122        return;
123    }
124
125    ggml_cuda_pdl_sync();
126
127    const int64_t i0 = ir % n_embd;
128    const int64_t it = ir / n_embd;
129
130    float sum = x[i0*sx0 + it*sx2] * weights[it*sw1];
131    for (int64_t ih = 1; ih < hc; ++ih) {
132        const float xv = x[i0*sx0 + ih*sx1 + it*sx2];
133        const float wv = weights[ih*sw0 + it*sw1];
134        sum += xv * wv;
135    }
136
137    dst[i0*sd0 + it*sd1] = sum;
138}
139
140static __global__ void dsv4_hc_post_f32(
141        const float * x,
142        const float * residual,
143        const float * post,
144        const float * comb,
145        float * dst,
146        int64_t n_embd,
147        int64_t hc,
148        int64_t n_tokens,
149        int64_t sx0,
150        int64_t sx1,
151        int64_t sr0,
152        int64_t sr1,
153        int64_t sr2,
154        int64_t sp0,
155        int64_t sp1,
156        int64_t sc0,
157        int64_t sc1,
158        int64_t sc2,
159        int64_t sd0,
160        int64_t sd1,
161        int64_t sd2) {
162    ggml_cuda_pdl_lc();
163    const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
164    const int64_t nr = n_embd * hc * n_tokens;
165
166    if (ir >= nr) {
167        return;
168    }
169
170    ggml_cuda_pdl_sync();
171
172    const int64_t i0   = ir % n_embd;
173    const int64_t idst = (ir / n_embd) % hc;
174    const int64_t it   = ir / (n_embd * hc);
175
176    float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1];
177    for (int64_t isrc = 0; isrc < hc; ++isrc) {
178        sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2];
179    }
180
181    dst[i0*sd0 + idst*sd1 + it*sd2] = sum;
182}
183
184void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
185    const ggml_tensor * mixes = dst->src[0];
186    const ggml_tensor * scale = dst->src[1];
187    const ggml_tensor * base  = dst->src[2];
188
189    GGML_ASSERT(mixes->type == GGML_TYPE_F32);
190    GGML_ASSERT(scale->type == GGML_TYPE_F32);
191    GGML_ASSERT(base->type == GGML_TYPE_F32);
192    GGML_ASSERT(dst->type == GGML_TYPE_F32);
193
194    constexpr int64_t hc_mix_dim = (2 + DSV4_HC)*DSV4_HC;
195
196    GGML_ASSERT(mixes->ne[0] == hc_mix_dim);
197    GGML_ASSERT(dst->ne[0] == DSV4_HC);
198    GGML_ASSERT(dst->ne[1] == DSV4_HC);
199    GGML_ASSERT(dst->ne[2] == mixes->ne[1]);
200    GGML_ASSERT(scale->ne[0] >= 3);
201    GGML_ASSERT(base->ne[0] == hc_mix_dim);
202
203    GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb);
204    GGML_TENSOR_LOCALS(size_t, nbs, scale, nb);
205    GGML_TENSOR_LOCALS(size_t, nbb, base,  nb);
206    GGML_TENSOR_LOCALS(size_t, nbd, dst,   nb);
207
208    const int64_t n_tokens = mixes->ne[1];
209    const float eps = ggml_get_op_params_f32(dst, 0);
210    const int32_t n_iter = ggml_get_op_params_i32(dst, 1);
211
212    const int block_size = 256;
213    const dim3 block_dims(block_size, 1, 1);
214    const dim3 grid_dims((n_tokens + block_size - 1) / block_size, 1, 1);
215    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
216
217    ggml_cuda_kernel_launch(dsv4_hc_comb_f32, launch_params,
218            (const float *) mixes->data, (const float *) scale->data, (const float *) base->data, (float *) dst->data,
219            n_tokens,
220            nbm0 / sizeof(float), nbm1 / sizeof(float),
221            nbs0 / sizeof(float),
222            nbb0 / sizeof(float),
223            nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float),
224            eps, n_iter);
225}
226
227void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
228    const ggml_tensor * x       = dst->src[0];
229    const ggml_tensor * weights = dst->src[1];
230
231    GGML_ASSERT(x->type == GGML_TYPE_F32);
232    GGML_ASSERT(weights->type == GGML_TYPE_F32);
233    GGML_ASSERT(dst->type == GGML_TYPE_F32);
234
235    GGML_TENSOR_LOCALS(size_t, nbx, x,       nb);
236    GGML_TENSOR_LOCALS(size_t, nbw, weights, nb);
237    GGML_TENSOR_LOCALS(size_t, nbd, dst,     nb);
238
239    const int64_t n_embd   = x->ne[0];
240    const int64_t hc       = x->ne[1];
241    const int64_t n_tokens = x->ne[2];
242
243    const int block_size = 256;
244    const int64_t nr = n_embd * n_tokens;
245    const dim3 block_dims(block_size, 1, 1);
246    const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1);
247    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
248
249    ggml_cuda_kernel_launch(dsv4_hc_pre_f32, launch_params,
250            (const float *) x->data, (const float *) weights->data, (float *) dst->data,
251            n_embd, hc, n_tokens,
252            nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float),
253            nbw0 / sizeof(float), nbw1 / sizeof(float),
254            nbd0 / sizeof(float), nbd1 / sizeof(float));
255}
256
257void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
258    const ggml_tensor * x        = dst->src[0];
259    const ggml_tensor * residual = dst->src[1];
260    const ggml_tensor * post     = dst->src[2];
261    const ggml_tensor * comb     = dst->src[3];
262
263    GGML_ASSERT(x->type == GGML_TYPE_F32);
264    GGML_ASSERT(residual->type == GGML_TYPE_F32);
265    GGML_ASSERT(post->type == GGML_TYPE_F32);
266    GGML_ASSERT(comb->type == GGML_TYPE_F32);
267    GGML_ASSERT(dst->type == GGML_TYPE_F32);
268
269    GGML_TENSOR_LOCALS(size_t, nbx, x,        nb);
270    GGML_TENSOR_LOCALS(size_t, nbr, residual, nb);
271    GGML_TENSOR_LOCALS(size_t, nbp, post,     nb);
272    GGML_TENSOR_LOCALS(size_t, nbc, comb,     nb);
273    GGML_TENSOR_LOCALS(size_t, nbd, dst,      nb);
274
275    const int64_t n_embd   = x->ne[0];
276    const int64_t n_tokens = x->ne[1];
277    const int64_t hc       = residual->ne[1];
278
279    const int block_size = 256;
280    const int64_t nr = n_embd * hc * n_tokens;
281    const dim3 block_dims(block_size, 1, 1);
282    const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1);
283    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
284
285    ggml_cuda_kernel_launch(dsv4_hc_post_f32, launch_params,
286            (const float *) x->data, (const float *) residual->data,
287            (const float *) post->data, (const float *) comb->data, (float *) dst->data,
288            n_embd, hc, n_tokens,
289            nbx0 / sizeof(float), nbx1 / sizeof(float),
290            nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float),
291            nbp0 / sizeof(float), nbp1 / sizeof(float),
292            nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float),
293            nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float));
294}