char-fork/llama.cpp

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

Daniel Beveniusmtmd : add Nemotron 3 Nano Omni support (parakeet) (#22520)7e1e28cae

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
19.2 KiB421 linesraw
1#include "models.h"
2
3static constexpr int PARAKEET_LOCAL_ATTN_THRESHOLD = 8192;
4static constexpr int PARAKEET_LOCAL_ATTN_WINDOW    = 128;
5
6// conv subsampling + conformer encoder
7ggml_cgraph * clip_graph_parakeet::build() {
8
9    // Conv subsampling
10    ggml_tensor * inp = build_inp_raw(1);
11    inp = ggml_cont(ctx0, ggml_transpose(ctx0, inp));
12
13    // [freq, time, channels, batch]
14    ggml_tensor * cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[0], inp, 2, 2, 1, 1, 1, 1);
15    cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[0]);
16    cb(cur, "pre_conv_0", -1);
17
18    cur = ggml_relu(ctx0, cur);
19    cb(cur, "pre_conv_0_relu", -1);
20
21    // [freq, time, channels, batch]
22    cur = ggml_conv_2d_dw_direct(ctx0, model.pre_encode_conv_X_w[2], cur, 2, 2, 1, 1, 1, 1);
23    cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[2]);
24    cb(cur, "pre_conv_2", -1);
25
26    // [freq, time, channels, batch]
27    cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[3], cur, 1, 1, 0, 0, 1, 1);
28    cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[3]);
29    cb(cur, "pre_conv_3", -1);
30
31    cur = ggml_relu(ctx0, cur);
32    cb(cur, "pre_conv_3_relu", -1);
33
34    // [freq, time, channels, batch]
35    cur = ggml_conv_2d_dw_direct(ctx0, model.pre_encode_conv_X_w[5], cur, 2, 2, 1, 1, 1, 1);
36    cb(cur, "pre_conv_5_direct", -1);
37    cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[5]);
38    cb(cur, "pre_conv_5", -1);
39
40    // [freq, time, channels, batch]
41    cur = ggml_conv_2d(ctx0, model.pre_encode_conv_X_w[6], cur, 1, 1, 0, 0, 1, 1);
42    cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[6]);
43    cb(cur, "pre_conv_6", -1);
44
45    cur = ggml_relu(ctx0, cur);
46    cb(cur, "pre_conv_6_relu", -1);
47
48    // [freq, time, chan]
49    cur = ggml_permute(ctx0, cur, 0, 2, 1, 3);
50    // [freq, chan, time]
51    cur = ggml_cont(ctx0, cur);
52
53    const int n_freq   = cur->ne[0];
54    const int n_chan   = cur->ne[1];
55    const int n_frames = cur->ne[2];
56
57    // [freq, time, chan, batch] -> [(freq * chan), time]
58    cur = ggml_reshape_2d(ctx0, cur, n_freq * n_chan, n_frames);
59
60    cur = build_mm(model.pre_encode_out_w, cur);
61    cur = ggml_add(ctx0, cur, model.pre_encode_out_b);
62
63    ggml_set_name(cur, "pre_enc_out");
64
65    // Encoder
66
67    const auto & hparams  = model.hparams;
68    const int n_layer     = hparams.n_layer;
69    const int n_state     = hparams.n_embd;
70    const float fc_factor = 0.5f;
71
72    const int  n_time      = cur->ne[1];
73    const bool local_attn  = n_time > PARAKEET_LOCAL_ATTN_THRESHOLD;
74    const int  att_left    = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1;
75    const int  att_right   = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1;
76    const int  window_size = local_attn ? att_left + att_right + 1 : 2 * n_time - 1;
77    const int  d_half      = n_state / 2;
78    const int  mask_dim    = local_attn ? window_size : n_time;
79
80    // mask [key, n_time]
81    struct ggml_tensor * attn_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, mask_dim, n_time);
82    ggml_set_name(attn_mask, "attn_mask");
83    ggml_set_input(attn_mask);
84
85    struct ggml_tensor * local_mask = nullptr;
86    if (local_attn) {
87        const int chunk = att_left + att_right;
88        local_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, chunk + window_size - 1, chunk);
89        ggml_set_name(local_mask, "local_mask");
90        ggml_set_input(local_mask);
91    }
92
93    struct ggml_tensor * pos_freqs = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, d_half);
94    ggml_set_name(pos_freqs, "pos_freqs");
95    ggml_set_input(pos_freqs);
96
97    struct ggml_tensor * rel_positions = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, window_size);
98    ggml_set_name(rel_positions, "rel_positions");
99    ggml_set_input(rel_positions);
100
101    struct ggml_tensor * freqs = ggml_repeat_4d(ctx0, pos_freqs, d_half, window_size, 1, 1);
102    struct ggml_tensor * theta = ggml_mul(ctx0, freqs, rel_positions);
103
104    struct ggml_tensor * sin = ggml_reshape_3d(ctx0, ggml_sin(ctx0, theta), 1, d_half, window_size);
105    struct ggml_tensor * cos = ggml_reshape_3d(ctx0, ggml_cos(ctx0, theta), 1, d_half, window_size);
106    struct ggml_tensor * pos_emb = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_concat(ctx0, sin, cos, 0)), n_state, window_size);
107    ggml_set_name(pos_emb, "pos_emb");
108
109    for (int il = 0; il < n_layer; ++il) {
110        const auto & layer = model.layers[il];
111        // FFN1
112        {
113            struct ggml_tensor * residual = cur;
114            ggml_format_name(cur, "enc_%d_res", il);
115
116            // norm
117            cur = ggml_norm(ctx0, cur, hparams.eps);
118            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ff_norm_w), layer.ff_norm_b);
119            ggml_format_name(cur, "enc_%d_ffn_norm_1", il);
120
121            cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_SILU, il);
122            ggml_format_name(cur, "enc_%d_ffn_1", il);
123
124            cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, fc_factor));
125            ggml_format_name(cur, "enc_%d_res_ffn", il);
126        }
127
128        // self attention block using relative positional encoding from model.position_embedding.
129        {
130            // [feat, time_frames, 1, 1]
131            struct ggml_tensor * residual = cur;
132
133            cur = ggml_norm(ctx0, cur, hparams.eps);
134            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ln_1_w), layer.ln_1_b);
135            ggml_format_name(cur, "enc_%d_attn_norm", il);
136
137            const int n_head = hparams.n_head;
138            const int d_head = n_state / n_head;
139
140            // [feat, time_frames, 1, 1]
141            struct ggml_tensor * Q_cur = build_mm(layer.q_w, cur);
142            struct ggml_tensor * K_cur = build_mm(layer.k_w, cur);
143            struct ggml_tensor * V_cur = build_mm(layer.v_w, cur);
144
145            // [d_head, n_heads, n_time, 1]
146            Q_cur = ggml_reshape_3d(ctx0, Q_cur, d_head, n_head, n_time);
147            K_cur = ggml_reshape_3d(ctx0, K_cur, d_head, n_head, n_time);
148            V_cur = ggml_reshape_3d(ctx0, V_cur, d_head, n_head, n_time);
149
150            // [n_state, window_size]
151            struct ggml_tensor * pos = build_mm(layer.linear_pos_w, pos_emb);
152            // [feat, head, window_size, 1]
153            pos = ggml_reshape_3d(ctx0, pos, d_head, n_head, pos_emb->ne[1]);
154            // [feat, window_size, head, 1]
155            pos = ggml_cont(ctx0, ggml_permute(ctx0, pos, 0, 2, 1, 3));
156            ggml_format_name(pos, "enc_%d_attn_pos", il);
157
158            if (local_attn) {
159                const int  chunk         = att_left + att_right;
160                const int  n_group       = (n_time + chunk - 1) / chunk;
161                const int  n_time_padded = n_group * chunk;
162                const int  n_kv_chunk    = chunk + window_size - 1;
163                const int  n_kv_dense    = n_kv_chunk * n_group;
164                const bool need_padding  = n_time_padded > n_time;
165
166                Q_cur = ggml_cont(ctx0, ggml_permute(ctx0, Q_cur, 0, 2, 1, 3));
167                K_cur = ggml_cont(ctx0, ggml_permute(ctx0, K_cur, 0, 2, 1, 3));
168                V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 0, 2, 1, 3));
169
170                // content bias
171                struct ggml_tensor * bias_u = ggml_reshape_3d(ctx0, layer.pos_bias_u, d_head, 1, n_head);
172                struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, bias_u);
173
174                // position bias
175                struct ggml_tensor * bias_v = ggml_reshape_3d(ctx0, layer.pos_bias_v, d_head, 1, n_head);
176                struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, bias_v);
177
178                // right pad the time dimension
179                struct ggml_tensor * Q_u_padded = need_padding ?
180                    ggml_pad_ext(ctx0, Q_u, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : Q_u;
181                Q_u_padded = ggml_reshape_4d(ctx0, Q_u_padded, d_head, chunk, n_group, n_head);
182
183                // pad front and back for the first and last time frames
184                struct ggml_tensor * K_padded = ggml_pad_ext(ctx0, K_cur, 0, 0, att_left, att_right, 0, 0, 0, 0);
185                if (n_kv_dense > K_padded->ne[1]) {
186                    K_padded = ggml_pad_ext(ctx0, K_padded, 0, 0, 0, n_kv_dense - K_padded->ne[1], 0, 0, 0, 0);
187                }
188
189                // sliding window view: each group spans n_kv_chunk keys but steps by chunk
190                struct ggml_tensor * K_chunk = ggml_view_4d(ctx0, K_padded,
191                        d_head, n_kv_chunk, n_group, n_head,
192                        K_padded->nb[1],
193                        (size_t) chunk * K_padded->nb[1],
194                        K_padded->nb[2],
195                        0);
196                K_chunk = ggml_cont(ctx0, K_chunk);
197
198                struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_chunk, Q_u_padded);
199
200                // trim the dense output down to window_size scores per query
201                content_scores = ggml_view_4d(ctx0, content_scores,
202                        window_size, chunk, n_group, n_head,
203                        (size_t) (chunk + window_size) * content_scores->nb[0],
204                        content_scores->nb[2],
205                        content_scores->nb[3],
206                        0);
207                content_scores = ggml_cont(ctx0, content_scores);
208
209                // ungroup: [window_size, n_time_padded, n_head]
210                content_scores = ggml_reshape_3d(ctx0, content_scores, window_size, n_time_padded, n_head);
211                if (need_padding) {
212                    content_scores = ggml_view_3d(ctx0, content_scores,
213                            window_size, n_time, n_head,
214                            content_scores->nb[1],
215                            content_scores->nb[2],
216                            0);
217                }
218
219                // Q_v: [d_head, time, head]
220                Q_v = ggml_cont(ctx0, ggml_permute(ctx0, Q_v, 0, 2, 1, 3));
221                struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v);
222
223                struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores);
224                attn_scores = ggml_soft_max_ext(ctx0, attn_scores, attn_mask, 1.0f / std::sqrt(d_head), 0.0f);
225                ggml_format_name(attn_scores, "enc_%d_attn_probs", il);
226
227                // expand probs back to n_kv_chunk width for the V matmul
228                struct ggml_tensor * probs_padded = need_padding ?
229                    ggml_pad_ext(ctx0, attn_scores, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : attn_scores;
230
231                probs_padded = ggml_reshape_4d(ctx0, probs_padded, window_size, chunk, n_group, n_head);
232                probs_padded = ggml_pad_ext(ctx0, probs_padded, 0, chunk, 0, 0, 0, 0, 0, 0);
233                probs_padded = ggml_view_4d(ctx0, probs_padded,
234                        n_kv_chunk, chunk, n_group, n_head,
235                        (size_t) n_kv_chunk * probs_padded->nb[0],
236                        probs_padded->nb[2],
237                        probs_padded->nb[3],
238                        0);
239                probs_padded = ggml_cont(ctx0, probs_padded);
240                probs_padded = ggml_mul(ctx0, probs_padded, local_mask);
241
242                struct ggml_tensor * V_padded = ggml_pad_ext(ctx0, V_cur, 0, 0, att_left, att_right, 0, 0, 0, 0);
243                if (n_kv_dense > V_padded->ne[1]) {
244                    V_padded = ggml_pad_ext(ctx0, V_padded, 0, 0, 0, n_kv_dense - V_padded->ne[1], 0, 0, 0, 0);
245                }
246                V_padded = ggml_cont(ctx0, ggml_transpose(ctx0, V_padded));
247
248                struct ggml_tensor * V_chunk = ggml_view_4d(ctx0, V_padded,
249                        n_kv_chunk, d_head, n_group, n_head,
250                        V_padded->nb[1],
251                        (size_t) chunk * V_padded->nb[0],
252                        V_padded->nb[2],
253                        0);
254                V_chunk = ggml_cont(ctx0, V_chunk);
255
256                cur = ggml_mul_mat(ctx0, V_chunk, probs_padded);
257                cur = ggml_reshape_3d(ctx0, cur, d_head, n_time_padded, n_head);
258                if (need_padding) {
259                    cur = ggml_view_3d(ctx0, cur, d_head, n_time, n_head, cur->nb[1], cur->nb[2], 0);
260                }
261                cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3));
262                cur = ggml_reshape_2d(ctx0, cur, n_state, n_time);
263                cur = build_mm(layer.o_w, cur);
264            } else {
265                // full attention
266                struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, layer.pos_bias_u);
267                ggml_format_name(Q_u, "enc_%d_attn_q_u", il);
268
269                struct ggml_tensor * K_prep = ggml_permute(ctx0, K_cur, 0, 2, 1, 3);
270                struct ggml_tensor * Q_prep = ggml_permute(ctx0, Q_u,   0, 2, 1, 3);
271                struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_prep, Q_prep);
272                ggml_format_name(content_scores, "enc_%d_attn_content_scores", il);
273
274                struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, layer.pos_bias_v);
275                ggml_format_name(Q_v, "enc_%d_attn_q_v", il);
276
277                Q_v = ggml_permute(ctx0, Q_v, 0, 2, 1, 3);
278                Q_v = ggml_cont(ctx0, Q_v);
279                ggml_format_name(Q_v, "enc_%d_attn_q_v_perm", il);
280
281                struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v);
282                ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos", il);
283
284                // Relative positional shift
285                {
286                    const auto pos_window = rel_pos_scores->ne[0];
287                    const auto n_frame    = rel_pos_scores->ne[1];
288                    const auto n_head     = rel_pos_scores->ne[2];
289
290                    rel_pos_scores = ggml_pad(ctx0, rel_pos_scores, 1, 0, 0, 0);
291                    rel_pos_scores = ggml_roll(ctx0, rel_pos_scores, 1, 0, 0, 0);
292
293                    rel_pos_scores = ggml_reshape_3d(ctx0, rel_pos_scores, n_frame, pos_window + 1, n_head);
294                    rel_pos_scores = ggml_cont(ctx0, rel_pos_scores);
295                    ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_reshaped", il);
296
297                    int center = pos_window / 2;
298                    size_t offset = rel_pos_scores->nb[0] * (center+1);
299
300                    rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores,
301                                                  n_frame, pos_window, n_head,
302                                                  (pos_window) * 4,
303                                                  rel_pos_scores->nb[2],
304                                                  offset);
305                    rel_pos_scores = ggml_cont(ctx0, rel_pos_scores);
306                    ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted", il);
307
308                    rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores,
309                                                  content_scores->ne[0],
310                                                  content_scores->ne[1],
311                                                  rel_pos_scores->ne[2],
312                                                  rel_pos_scores->nb[1],
313                                                  rel_pos_scores->nb[2],
314                                                  0);
315                    rel_pos_scores = ggml_cont(ctx0, rel_pos_scores);
316                    ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted_view", il);
317                }
318
319                struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores);
320                ggml_format_name(attn_scores, "enc_%d_attn_scores", il);
321                attn_scores = ggml_scale(ctx0, attn_scores, 1.0f / std::sqrt(d_head));
322                attn_scores = ggml_add(ctx0, attn_scores, attn_mask);
323                ggml_format_name(attn_scores, "enc_%d_attn_scores_scaled", il);
324
325                struct ggml_tensor * probs = ggml_soft_max(ctx0, attn_scores);
326                ggml_format_name(probs, "enc_%d_attn_probs", il);
327
328                V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 1, 2, 0, 3));
329                ggml_format_name(V_cur, "enc_%d_attn_v_cur", il);
330                cur = ggml_mul_mat(ctx0, probs, V_cur);
331                ggml_format_name(cur, "enc_%d_attn_inp", il);
332
333                cur = ggml_permute(ctx0, cur, 2, 0, 1, 3);
334                cur = ggml_cont_2d(ctx0, cur, n_state, n_time);
335                cur = build_mm(layer.o_w, cur);
336            }
337            ggml_format_name(cur, "enc_%d_attn_out", il);
338
339            cur = ggml_add(ctx0, residual, cur);
340            ggml_format_name(cur, "enc_%d_attn_res", il);
341        }
342
343        // Convolution
344        {
345            struct ggml_tensor * residual = cur;
346            ggml_format_name(cur, "enc_%d_residual_conv", il);
347
348            cur = ggml_norm(ctx0, cur, hparams.eps);
349            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_conv_w), layer.norm_conv_b);
350            ggml_format_name(cur, "enc_%d_norm_conv", il);
351
352            // pointwise 1d convolution:
353            cur = build_mm(layer.conv_pw1_w, cur);
354            ggml_format_name(cur, "enc_%d_conv_pw1", il);
355
356            {
357                int64_t d = cur->ne[0] / 2;
358                struct ggml_tensor * signal = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], 0);
359                struct ggml_tensor * gate   = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], d * cur->nb[0]);
360
361                cur = ggml_mul(ctx0, signal, ggml_sigmoid(ctx0, gate));
362                ggml_format_name(cur, "enc_%d_conv_glu", il);
363            }
364
365            cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
366
367            // use ggml_ssm_conv for f32 precision
368            const int dw_pad = (hparams.audio_conv_kernel_size - 1) / 2;
369            cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0);
370            cur = ggml_roll(ctx0, cur, dw_pad, 0, 0, 0);
371            cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0);
372            ggml_format_name(cur, "enc_%d_conv_dw_pad", il);
373
374            cur = ggml_ssm_conv(ctx0, cur, layer.conv_dw_w);
375            ggml_format_name(cur, "enc_%d_conv_1d_dw", il);
376
377            cur = ggml_sub(ctx0, cur, layer.conv_norm_mean);
378            struct ggml_tensor * std = ggml_sqrt(ctx0, layer.conv_norm_var);
379            cur = ggml_div(ctx0, cur, std);
380            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.conv_norm_w), layer.conv_norm_b);
381            ggml_format_name(cur, "enc_%d_conv_bn", il);
382
383            cur = ggml_silu(ctx0, cur);
384            ggml_format_name(cur, "enc_%d_conv_silu", il);
385
386            cur = build_mm(layer.conv_pw2_w, cur);
387            ggml_format_name(cur, "enc_%d_conv_pw2", il);
388
389            cur = ggml_add(ctx0, residual, cur);
390            ggml_format_name(cur, "enc_%d_conv_res", il);
391        }
392
393        // FFN2
394        {
395            struct ggml_tensor * residual = cur;
396            cur = ggml_norm(ctx0, cur, hparams.eps);
397            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ff_norm_1_w), layer.ff_norm_1_b);
398            ggml_format_name(cur, "enc_%d_ffn_norm_2", il);
399
400            cur = build_ffn(cur, layer.ff_up_1_w, nullptr, nullptr, nullptr, layer.ff_down_1_w, nullptr, FFN_SILU, il);
401            cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, 0.5));
402            ggml_format_name(cur, "enc_%d_ffn_res", il);
403        }
404
405        cur = ggml_norm(ctx0, cur, hparams.eps);
406        cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.ln_2_w), layer.ln_2_b);
407    }
408
409    cb(cur, "encoder_out", -1);
410
411    cur = ggml_rms_norm(ctx0, cur, 1e-6);
412    cur = ggml_mul(ctx0, cur, model.mm_norm_pre_w);
413    cb(cur, "sound_projection.norm", -1);
414
415    cur = build_ffn(cur, model.mm_0_w, model.mm_0_b, nullptr, nullptr, model.mm_1_w, model.mm_1_b, FFN_RELU_SQR, -1);
416    cb(cur, "projected", -1);
417
418    ggml_build_forward_expand(gf, cur);
419
420    return gf;
421}