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
7.9 KiB164 linesraw
1#include "models.h"
2
3#include "llama-memory-recurrent.h"
4
5llm_build_rwkv6_base::llm_build_rwkv6_base(const llama_model & model, const llm_graph_params & params) :
6    llm_graph_context(params),
7    model(model) {}
8
9ggml_tensor * llm_build_rwkv6_base::build_rwkv6_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_RWKV6:
16            {
17                ggml_tensor * xk = ggml_add(ctx0, ggml_mul(ctx0, sx, layer->channel_mix_lerp_k), cur);
18                ggml_tensor * xr = ggml_add(ctx0, ggml_mul(ctx0, sx, layer->channel_mix_lerp_r), cur);
19
20                ggml_tensor * r = ggml_sigmoid(ctx0, build_lora_mm(layer->channel_mix_receptance, xr));
21                ggml_tensor * k = ggml_sqr(ctx0, ggml_relu(ctx0, build_lora_mm(layer->channel_mix_key, xk)));
22                cur             = ggml_mul(ctx0, r, build_lora_mm(layer->channel_mix_value, k));
23            }
24            break;
25        default:
26            GGML_ABORT("fatal error");
27    }
28    return cur;
29}
30
31ggml_tensor * llm_build_rwkv6_base::build_rwkv6_time_mix(llm_graph_input_rs * inp,
32                                                         ggml_tensor *        cur,
33                                                         ggml_tensor *        x_prev,
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_seq_tokens = ubatch.n_seq_tokens;
41    const auto n_embd       = hparams.n_embd;
42    const auto head_size    = hparams.wkv_head_size;
43    const auto n_head       = n_embd / head_size;
44    const auto n_head_kv    = hparams.n_head_kv(il);
45
46    const auto kv_head = mctx_cur->get_head();
47
48    const auto & layer = model.layers[il];
49
50    bool is_qrwkv = layer.time_mix_first == nullptr;
51
52    ggml_tensor * sx = ggml_sub(ctx0, x_prev, cur);
53
54    sx  = ggml_reshape_2d(ctx0, sx, n_embd, n_tokens);
55    cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens);
56
57    ggml_tensor * xxx = ggml_add(ctx0, ggml_mul(ctx0, sx, layer.time_mix_lerp_x), cur);
58
59    xxx = ggml_reshape_4d(ctx0, ggml_tanh(ctx0, ggml_mul_mat(ctx0, layer.time_mix_w1, xxx)),
60                          layer.time_mix_w1->ne[1] / 5, 1, 5, n_tokens);
61
62    xxx = ggml_cont(ctx0, ggml_permute(ctx0, xxx, 0, 1, 3, 2));
63
64    xxx = ggml_mul_mat(
65        ctx0, ggml_reshape_4d(ctx0, layer.time_mix_w2, layer.time_mix_w2->ne[0], layer.time_mix_w2->ne[1], 1, 5), xxx);
66
67    ggml_tensor *xw, *xk, *xv, *xr, *xg;
68    if (layer.time_mix_lerp_fused) {
69        // fusing these weights makes some performance improvement
70        sx  = ggml_reshape_3d(ctx0, sx, n_embd, 1, n_tokens);
71        cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
72        xxx = ggml_add(ctx0, ggml_mul(ctx0, ggml_add(ctx0, xxx, layer.time_mix_lerp_fused), sx), cur);
73        xw  = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], 0);
74        xk  = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * sizeof(float));
75        xv  = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 2 * sizeof(float));
76        xr  = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 3 * sizeof(float));
77        xg  = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 4 * sizeof(float));
78    } else {
79        // for backward compatibility
80        xw = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], 0);
81        xk = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * sizeof(float));
82        xv = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 2 * sizeof(float));
83        xr = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 3 * sizeof(float));
84        xg = ggml_view_2d(ctx0, xxx, n_embd, n_tokens, xxx->nb[1], n_embd * n_tokens * 4 * sizeof(float));
85
86        xw = ggml_add(ctx0, ggml_mul(ctx0, ggml_add(ctx0, xw, layer.time_mix_lerp_w), sx), cur);
87        xk = ggml_add(ctx0, ggml_mul(ctx0, ggml_add(ctx0, xk, layer.time_mix_lerp_k), sx), cur);
88        xv = ggml_add(ctx0, ggml_mul(ctx0, ggml_add(ctx0, xv, layer.time_mix_lerp_v), sx), cur);
89        xr = ggml_add(ctx0, ggml_mul(ctx0, ggml_add(ctx0, xr, layer.time_mix_lerp_r), sx), cur);
90        xg = ggml_add(ctx0, ggml_mul(ctx0, ggml_add(ctx0, xg, layer.time_mix_lerp_g), sx), cur);
91    }
92    ggml_tensor * r = build_lora_mm(layer.time_mix_receptance, xr);
93    ggml_tensor * k = build_lora_mm(layer.time_mix_key, xk);
94    ggml_tensor * v = build_lora_mm(layer.time_mix_value, xv);
95    if (layer.time_mix_receptance_b) {
96        r = ggml_add(ctx0, r, layer.time_mix_receptance_b);
97    }
98    if (layer.time_mix_key_b) {
99        k = ggml_add(ctx0, k, layer.time_mix_key_b);
100    }
101    if (layer.time_mix_value_b) {
102        v = ggml_add(ctx0, v, layer.time_mix_value_b);
103    }
104    ggml_tensor * g = build_lora_mm(layer.time_mix_gate, xg);
105    if (is_qrwkv) {
106        g = ggml_sigmoid(ctx0, g);
107    } else {
108        g = ggml_silu(ctx0, g);
109    }
110    if (n_head_kv != 0 && n_head_kv != n_head) {
111        GGML_ASSERT(n_head % n_head_kv == 0);
112        k                 = ggml_reshape_4d(ctx0, k, head_size, 1, n_head_kv, n_tokens);
113        v                 = ggml_reshape_4d(ctx0, v, head_size, 1, n_head_kv, n_tokens);
114        ggml_tensor * tmp = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, head_size, n_head / n_head_kv, n_head_kv, n_tokens);
115        k                 = ggml_repeat(ctx0, k, tmp);
116        v                 = ggml_repeat(ctx0, v, tmp);
117    }
118    k = ggml_reshape_3d(ctx0, k, head_size, n_head, n_tokens);
119    v = ggml_reshape_3d(ctx0, v, head_size, n_head, n_tokens);
120    r = ggml_reshape_3d(ctx0, r, head_size, n_head, n_tokens);
121
122    ggml_tensor * w =
123        ggml_mul_mat(ctx0, layer.time_mix_decay_w2, ggml_tanh(ctx0, ggml_mul_mat(ctx0, layer.time_mix_decay_w1, xw)));
124
125    w = ggml_add(ctx0, w, layer.time_mix_decay);
126    w = ggml_exp(ctx0, ggml_neg(ctx0, ggml_exp(ctx0, w)));
127    w = ggml_reshape_3d(ctx0, w, head_size, n_head, n_tokens);
128
129    if (is_qrwkv) {
130        // k = k * (1 - w)
131        k = ggml_sub(ctx0, k, ggml_mul(ctx0, k, w));
132    }
133    ggml_tensor * wkv_state = build_rs(inp, mctx_cur->get_s_l(il), hparams.n_embd_s(), n_seqs);
134
135    ggml_tensor * wkv_output;
136    if (is_qrwkv) {
137        wkv_output = ggml_gated_linear_attn(ctx0, k, v, r, w, wkv_state, pow(head_size, -0.5f));
138    } else {
139        wkv_output = ggml_rwkv_wkv6(ctx0, k, v, r, layer.time_mix_first, w, wkv_state);
140    }
141    cur       = ggml_view_1d(ctx0, wkv_output, n_embd * n_tokens, 0);
142    wkv_state = ggml_view_1d(ctx0, wkv_output, n_embd * head_size * n_seqs, n_embd * n_tokens * sizeof(float));
143
144    ggml_build_forward_expand(
145        gf, ggml_cpy(ctx0, wkv_state,
146                     ggml_view_1d(ctx0, mctx_cur->get_s_l(il), hparams.n_embd_s() * n_seqs,
147                                  hparams.n_embd_s() * kv_head * ggml_element_size(mctx_cur->get_s_l(il)))));
148
149    if (!is_qrwkv) {
150        // group norm with head_count groups
151        cur = ggml_reshape_3d(ctx0, cur, n_embd / n_head, n_head, n_tokens);
152        cur = ggml_norm(ctx0, cur, 64e-5f);
153
154        // Convert back to regular vectors.
155        cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens);
156        cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.time_mix_ln), layer.time_mix_ln_b);
157    } else {
158        cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens);
159    }
160    cur = ggml_mul(ctx0, cur, g);
161    cur = build_lora_mm(layer.time_mix_output, cur);
162
163    return ggml_reshape_3d(ctx0, cur, n_embd, n_seq_tokens, n_seqs);
164}