char-fork/llama.cpp

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

Georgi Gerganovmodels : deduplicate delta-net graphs for Qwen family (#19597)cc45f2ada

char/aliceai-t5-moe
6.8 KiB137 linesraw
1#include "models.h"
2
3#include "llama-memory-recurrent.h"
4
5llm_build_rwkv7_base::llm_build_rwkv7_base(const llama_model & model, const llm_graph_params & params) :
6    llm_graph_context(params),
7    model(model) {}
8
9ggml_tensor * llm_build_rwkv7_base::build_rwkv7_channel_mix(const llama_layer * layer,
10                                                            ggml_tensor *       cur,
11                                                            ggml_tensor *       x_prev,
12                                                            llm_arch            arch) const {
13    ggml_tensor * sx = ggml_sub(ctx0, x_prev, cur);
14    switch (arch) {
15        case LLM_ARCH_RWKV7:
16            {
17                ggml_tensor * xk = ggml_add(ctx0, ggml_mul(ctx0, sx, layer->channel_mix_lerp_k), cur);
18
19                ggml_tensor * k = ggml_sqr(ctx0, ggml_relu(ctx0, build_lora_mm(layer->channel_mix_key, xk)));
20
21                cur = build_lora_mm(layer->channel_mix_value, k);
22            }
23            break;
24        default:
25            GGML_ABORT("fatal error");
26    }
27    return cur;
28}
29
30ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix(llm_graph_input_rs * inp,
31                                                         ggml_tensor *        cur,
32                                                         ggml_tensor *        x_prev,
33                                                         ggml_tensor *&       first_layer_value,
34                                                         const llama_ubatch & ubatch,
35                                                         int                  il) const {
36    const auto * mctx_cur = static_cast<const llama_memory_recurrent_context *>(mctx);
37
38    const auto n_tokens     = ubatch.n_tokens;
39    const auto n_seqs       = ubatch.n_seqs;
40    const auto n_embd       = hparams.n_embd;
41    const auto head_size    = hparams.wkv_head_size;
42    const auto head_count   = n_embd / head_size;
43    const auto n_seq_tokens = ubatch.n_seq_tokens;
44
45    const auto kv_head = mctx_cur->get_head();
46
47    const auto & layer = model.layers[il];
48
49    bool has_gating = layer.time_mix_g1 && layer.time_mix_g2;
50
51    ggml_tensor * sx    = ggml_sub(ctx0, x_prev, cur);
52    ggml_tensor * dummy = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_embd, n_seq_tokens, n_seqs, has_gating ? 6 : 5);
53    sx                  = ggml_repeat(ctx0, sx, dummy);
54
55    ggml_tensor * xxx = ggml_add(ctx0, ggml_mul(ctx0, sx, layer.time_mix_lerp_fused), cur);
56
57    ggml_tensor * xr = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], 0);
58    ggml_tensor * xw = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * sizeof(float));
59    ggml_tensor * xk = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 2 * sizeof(float));
60    ggml_tensor * xv = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 3 * sizeof(float));
61    ggml_tensor * xa = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 4 * sizeof(float));
62    ggml_tensor * xg =
63        has_gating ? ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 5 * sizeof(float)) :
64                     nullptr;
65
66    ggml_tensor * r = build_lora_mm(layer.time_mix_receptance, xr);
67    ggml_tensor * w = ggml_add(
68        ctx0, ggml_mul_mat(ctx0, layer.time_mix_w2, ggml_tanh(ctx0, ggml_mul_mat(ctx0, layer.time_mix_w1, xw))),
69        layer.time_mix_w0);
70    w = ggml_exp(ctx0, ggml_scale(ctx0, ggml_sigmoid(ctx0, w), -0.606531));
71
72    ggml_tensor * k = build_lora_mm(layer.time_mix_key, xk);
73    ggml_tensor * v = build_lora_mm(layer.time_mix_value, xv);
74    if (first_layer_value == nullptr) {
75        first_layer_value = v;
76    } else {
77        // Add the first layer value as a residual connection.
78        v = ggml_add(ctx0, v,
79                     ggml_mul(ctx0, ggml_sub(ctx0, first_layer_value, v),
80                              ggml_sigmoid(ctx0, ggml_add(ctx0,
81                                                          ggml_mul_mat(ctx0, layer.time_mix_v2,
82                                                                       ggml_mul_mat(ctx0, layer.time_mix_v1, xv)),
83                                                          layer.time_mix_v0))));
84    }
85    ggml_tensor * g = nullptr;
86    if (layer.time_mix_g1 && layer.time_mix_g2) {
87        g = ggml_mul_mat(ctx0, layer.time_mix_g2, ggml_sigmoid(ctx0, ggml_mul_mat(ctx0, layer.time_mix_g1, xg)));
88    }
89    ggml_tensor * a = ggml_sigmoid(
90        ctx0, ggml_add(ctx0, ggml_mul_mat(ctx0, layer.time_mix_a2, ggml_mul_mat(ctx0, layer.time_mix_a1, xa)),
91                       layer.time_mix_a0));
92
93    ggml_tensor * kk = ggml_reshape_3d(ctx0, ggml_mul(ctx0, k, layer.time_mix_k_k), head_size, head_count, n_tokens);
94    kk               = ggml_l2_norm(ctx0, kk, 1e-12);
95
96    ggml_tensor * ka = ggml_mul(ctx0, k, layer.time_mix_k_a);
97    k                = ggml_add(ctx0, k, ggml_sub(ctx0, ggml_mul(ctx0, a, ka), ka));
98
99    r = ggml_reshape_3d(ctx0, r, head_size, head_count, n_tokens);
100    w = ggml_reshape_3d(ctx0, w, head_size, head_count, n_tokens);
101    k = ggml_reshape_3d(ctx0, k, head_size, head_count, n_tokens);
102    v = ggml_reshape_3d(ctx0, v, head_size, head_count, n_tokens);
103    a = ggml_reshape_3d(ctx0, a, head_size, head_count, n_tokens);
104
105    ggml_tensor * wkv_state = build_rs(inp, mctx_cur->get_s_l(il), hparams.n_embd_s(), n_seqs);
106
107    ggml_tensor * wkv_output = ggml_rwkv_wkv7(ctx0, r, w, k, v, ggml_neg(ctx0, kk), ggml_mul(ctx0, kk, a), wkv_state);
108    cur                      = ggml_view_1d(ctx0, wkv_output, n_embd * n_tokens, 0);
109    wkv_state = ggml_view_1d(ctx0, wkv_output, n_embd * head_size * n_seqs, n_embd * n_tokens * sizeof(float));
110
111    ggml_build_forward_expand(
112        gf, ggml_cpy(ctx0, wkv_state,
113                     ggml_view_1d(ctx0, mctx_cur->get_s_l(il), hparams.n_embd_s() * n_seqs,
114                                  hparams.n_embd_s() * kv_head * ggml_element_size(mctx_cur->get_s_l(il)))));
115
116    if (layer.time_mix_ln && layer.time_mix_ln_b) {
117        // group norm with head_count groups
118        cur = ggml_reshape_3d(ctx0, cur, n_embd / head_count, head_count, n_tokens);
119        cur = ggml_norm(ctx0, cur, 64e-5f);
120
121        // Convert back to regular vectors.
122        cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens);
123        cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.time_mix_ln), layer.time_mix_ln_b);
124    } else {
125        cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens);
126    }
127    ggml_tensor * rk = ggml_sum_rows(
128        ctx0, ggml_mul(ctx0, ggml_mul(ctx0, k, r), ggml_reshape_2d(ctx0, layer.time_mix_r_k, head_size, head_count)));
129    cur = ggml_add(ctx0, cur, ggml_reshape_2d(ctx0, ggml_mul(ctx0, v, rk), n_embd, n_tokens));
130
131    if (has_gating) {
132        cur = ggml_mul(ctx0, cur, g);
133    }
134    cur = build_lora_mm(layer.time_mix_output, cur);
135
136    return ggml_reshape_3d(ctx0, cur, n_embd, n_seq_tokens, n_seqs);
137}