llama-cpp-sys-4 0.4.3

Low Level Bindings to llama.cpp
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218

#include "models.h"

ggml_cgraph * clip_graph_mimo_audio::build() {
    ggml_tensor * inp = build_inp_raw(1); // [n_frames, n_mel, 1]

    ggml_tensor * cur = ggml_conv_1d_ph(ctx0, model.conv1d_1_w, inp, 1, 1);
    cur = ggml_add(ctx0, cur, model.conv1d_1_b);
    cur = ggml_gelu_erf(ctx0, cur);

    cur = ggml_conv_1d_ph(ctx0, model.conv1d_2_w, cur, 2, 1);
    cur = ggml_add(ctx0, cur, model.conv1d_2_b);
    cur = ggml_gelu_erf(ctx0, cur);

    ggml_tensor * inpL = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [n_embd, n_pos]
    const int64_t n_pos = inpL->ne[1];
    cb(inpL, "after_conv1d", -1);

    GGML_ASSERT((int) hparams.wa_pattern_mode.size() == n_layer);

    ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos);
    ggml_set_name(inp_pos, "mimo_audio_positions");
    ggml_set_input(inp_pos);

    ggml_tensor * full_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, n_pos);
    ggml_set_name(full_mask, "mimo_audio_full_mask");
    ggml_set_input(full_mask);

    ggml_tensor * window_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, n_pos);
    ggml_set_name(window_mask, "mimo_audio_window_mask");
    ggml_set_input(window_mask);

    build_vit_opts opts;
    opts.attn_mask_layers.resize(n_layer);
    for (int il = 0; il < n_layer; il++) {
        opts.attn_mask_layers[il] = hparams.wa_pattern_mode[il] == -1 ? full_mask : window_mask;
    }
    // the skip connection below must be added before the post-transformer  norm,
    // so build_vit must not apply that norm itself
    opts.skip_post_ln = true;

    // encoder_skip_layer_id=3 (1-indexed) -> capture output of layer index 2
    const int skip_capture_il = 2;
    GGML_ASSERT(n_layer > skip_capture_il);
    ggml_tensor * skip_hidden = nullptr;
    opts.callback_layer_out = [&](ggml_tensor * layer_cur, int il) {
        if (il == skip_capture_il) {
            skip_hidden = layer_cur;
        }
    };

    auto add_pos = [&](ggml_tensor * x, const clip_layer &) {
        return ggml_rope_ext(ctx0, x, inp_pos, nullptr, d_head,
                             GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
    };

    inpL = build_vit(inpL, n_pos, NORM_TYPE_NORMAL, hparams.ffn_op, nullptr, add_pos, opts);
    inpL = ggml_reshape_2d(ctx0, inpL, n_embd, n_pos); // build_vit restores a (size-1) batch dim

    GGML_ASSERT(skip_hidden != nullptr);
    inpL = ggml_add(ctx0, inpL, skip_hidden);

    inpL = build_norm(inpL, model.post_ln_w, model.post_ln_b, NORM_TYPE_NORMAL, eps, -1);
    cb(inpL, "after_transformer", -1);

    // downsample: strided conv (no bias) + gelu + layernorm
    {
        ggml_tensor * ds = ggml_cont(ctx0, ggml_transpose(ctx0, inpL)); // [n_pos, n_embd]
        ds = ggml_conv_1d(ctx0, model.downsample_conv_w, ds, 2, 0, 1);
        ds = ggml_gelu_erf(ctx0, ds);
        ds = ggml_cont(ctx0, ggml_transpose(ctx0, ds)); // [n_embd, n_pos/2]
        ds = build_norm(ds, model.downsample_norm_w, model.downsample_norm_b, NORM_TYPE_NORMAL, eps, -1);
        inpL = ds;
    }
    cb(inpL, "after_downsample", -1);

    // RVQ quantize: codebook ne=[dim, max_bins, n_q]
    // quantize input vector to codes (type=I32)
    std::vector<ggml_tensor *> codes;
    {
        GGML_ASSERT(model.rvq_codebook != nullptr);
        const int64_t dim = model.rvq_codebook->ne[0];
        GGML_ASSERT(dim == inpL->ne[0]);
        GGML_ASSERT((int64_t) hparams.rvq_codebook_size.size() == model.rvq_codebook->ne[2]);

        ggml_tensor * residual = inpL; // [dim, n_pos_ds]

        for (size_t q = 0; q < hparams.rvq_codebook_size.size(); q++) {
            const int64_t bins = hparams.rvq_codebook_size[q];
            ggml_tensor * codebook_q = ggml_view_2d(ctx0, model.rvq_codebook, dim, bins,
                model.rvq_codebook->nb[1], q * model.rvq_codebook->nb[2]);
            codebook_q = ggml_cont(ctx0, codebook_q);

            ggml_tensor * codebook_norm = ggml_sum_rows(ctx0, ggml_sqr(ctx0, codebook_q)); // [1, bins]
            codebook_norm = ggml_cont(ctx0, ggml_transpose(ctx0, codebook_norm));          // [bins, 1]

            ggml_tensor * dot    = ggml_mul_mat(ctx0, codebook_q, residual); // [bins, n_pos_ds]
            ggml_tensor * scores = ggml_sub(ctx0, ggml_scale(ctx0, dot, 2.0f), codebook_norm);

            ggml_tensor * idx = ggml_argmax(ctx0, scores); // [n_pos_ds]
            codes.push_back(idx);

            ggml_tensor * quant = ggml_get_rows(ctx0, codebook_q, idx); // [dim, n_pos_ds]
            residual = ggml_sub(ctx0, residual, quant);
            cb(idx, "rvq_code", (int) q);
        }
    }

    // convert codes to LLM embeddings
    ggml_tensor * code_embd_sum = nullptr;
    {
        GGML_ASSERT(model.mm_a_code_embd != nullptr);
        const int64_t dim   = model.mm_a_code_embd->ne[0];
        const int64_t vocab = model.mm_a_code_embd->ne[1];
        GGML_ASSERT((int64_t) codes.size() == model.mm_a_code_embd->ne[2]);
        GGML_ASSERT(dim == inpL->ne[0]);

        for (size_t i = 0; i < codes.size(); i++) {
            ggml_tensor * table_i = ggml_view_2d(ctx0, model.mm_a_code_embd, dim, vocab,
                model.mm_a_code_embd->nb[1], i * model.mm_a_code_embd->nb[2]);
            table_i = ggml_cont(ctx0, table_i);

            ggml_tensor * embd_i = ggml_get_rows(ctx0, table_i, codes[i]); // [dim, n_pos_ds]
            code_embd_sum = code_embd_sum ? ggml_add(ctx0, code_embd_sum, embd_i) : embd_i;
        }
        cb(code_embd_sum, "code_embd_sum", -1);
    }

    // input_local_transformer
    // groups of `group_size` consecutive downsampled frames are processed together, attending only within their own group.
    // Implemented as a block-diagonal mask + in-group-repeating positions
    // (rather than a real batch dim) - same technique as the encoder's masks above, and as gemma4a's / deepseekocr2's chunked attention.

    // note: hand-rolled here instead of build_vit() because this is a second, independent layer stack
    // (own layer array/count, RMSNorm instead of LN, SiLU FFN, own RoPE theta)

    ggml_tensor * projected;
    {
        const int group_size = hparams.audio_local_group_size;
        GGML_ASSERT(group_size > 0);
        const int64_t n_pos_ds = code_embd_sum->ne[1];
        const int64_t n_groups = (n_pos_ds + group_size - 1) / group_size;
        const int64_t n_padded = n_groups * group_size;

        ggml_tensor * cur_local = code_embd_sum;
        if (n_padded != n_pos_ds) {
            cur_local = ggml_pad(ctx0, cur_local, 0, (int) (n_padded - n_pos_ds), 0, 0);
        }

        ggml_tensor * local_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_padded);
        ggml_set_name(local_pos, "mimo_audio_local_positions");
        ggml_set_input(local_pos);

        ggml_tensor * local_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_padded, n_padded);
        ggml_set_name(local_mask, "mimo_audio_local_mask");
        ggml_set_input(local_mask);

        const float local_rope_theta = 640000.0f; // audio_config.rope_theta (differs from the encoder's)
        auto apply_local_rope = [&](ggml_tensor * x) {
            return ggml_rope_ext(ctx0, x, local_pos, nullptr, d_head,
                                 GGML_ROPE_TYPE_NEOX, 0, local_rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
        };

        for (int il = 0; il < hparams.audio_local_n_layer; il++) {
            auto & layer = model.mm_a_local_layers[il];

            ggml_tensor * attn_in = build_norm(cur_local, layer.ln_1_w, nullptr, NORM_TYPE_RMS, eps, il);

            ggml_tensor * Qcur = build_mm(layer.q_w, attn_in);
            if (layer.q_b) {
                Qcur = ggml_add(ctx0, Qcur, layer.q_b);
            }
            ggml_tensor * Kcur = build_mm(layer.k_w, attn_in);
            if (layer.k_b) {
                Kcur = ggml_add(ctx0, Kcur, layer.k_b);
            }
            ggml_tensor * Vcur = build_mm(layer.v_w, attn_in);
            if (layer.v_b) {
                Vcur = ggml_add(ctx0, Vcur, layer.v_b);
            }

            Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_padded);
            Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_padded);
            Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_padded);

            Qcur = apply_local_rope(Qcur);
            Kcur = apply_local_rope(Kcur);

            ggml_tensor * attn_out = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, local_mask, kq_scale, il);
            cur_local = ggml_add(ctx0, cur_local, attn_out);

            ggml_tensor * ffn_in = build_norm(cur_local, layer.ln_2_w, nullptr, NORM_TYPE_RMS, eps, il);
            ggml_tensor * ffn_out = build_ffn(ffn_in,
                layer.ff_up_w, nullptr,
                layer.ff_gate_w, nullptr,
                layer.ff_down_w, nullptr,
                FFN_SILU, il);
            cur_local = ggml_add(ctx0, cur_local, ffn_out);
        }

        cur_local = build_norm(cur_local, model.mm_a_local_norm_w, nullptr, NORM_TYPE_RMS, eps, -1);
        cb(cur_local, "after_local_transformer", -1);

        // flatten each group of `group_size` frames into one (group_size*n_embd)-dim vector
        // (matching AudioProjection's flattened input)
        ggml_tensor * grouped = ggml_reshape_2d(ctx0, cur_local, n_embd * group_size, n_groups);

        // AudioProjection: Linear (no bias) -> GELU -> Linear (no bias)
        projected = build_ffn(grouped,
            model.mm_1_w, nullptr,
            nullptr, nullptr,
            model.mm_2_w, nullptr,
            FFN_GELU_ERF, -1);
        cb(projected, "after_projection", -1);
    }

    ggml_build_forward_expand(gf, projected);
    return gf;
}