1#include "models.h"
2
3llm_build_rwkv6::llm_build_rwkv6(const llama_model & model, const llm_graph_params & params) :
4 llm_build_rwkv6_base(model, params) {
5 GGML_ASSERT(hparams.token_shift_count == 2);
6
7 ggml_tensor * cur;
8 ggml_tensor * inpL;
9
10 inpL = build_inp_embd(tok_embd: model.tok_embd);
11 inpL = build_norm(cur: inpL, mw: model.tok_norm, mb: model.tok_norm_b, type: LLM_NORM, il: -1);
12
13 auto * rs_inp = build_rs_inp();
14
15 const auto n_embd = hparams.n_embd;
16 const auto n_seq_tokens = ubatch.n_seq_tokens;
17 const auto n_seqs = ubatch.n_seqs;
18
19 ggml_tensor * inp_out_ids = build_inp_out_ids();
20
21 for (int il = 0; il < n_layer; ++il) {
22 const llama_layer * layer = &model.layers[il];
23 inpL = ggml_reshape_3d(ctx: ctx0, a: inpL, ne0: n_embd, ne1: n_seq_tokens, ne2: n_seqs);
24
25 ggml_tensor * token_shift = build_rwkv_token_shift_load(inp: rs_inp, ubatch, il);
26
27 ggml_tensor * att_shift =
28 ggml_view_3d(ctx: ctx0, a: token_shift, ne0: n_embd, ne1: 1, ne2: n_seqs, nb1: token_shift->nb[1], nb2: token_shift->nb[2], offset: 0);
29 ggml_tensor * ffn_shift = ggml_view_3d(ctx: ctx0, a: token_shift, ne0: n_embd, ne1: 1, ne2: n_seqs, nb1: token_shift->nb[1],
30 nb2: token_shift->nb[2], offset: n_embd * ggml_element_size(tensor: token_shift));
31
32 ggml_tensor * att_norm = build_norm(cur: inpL, mw: layer->attn_norm, mb: layer->attn_norm_b, type: LLM_NORM, il);
33 cb(cur: att_norm, name: "attn_norm", il);
34
35 ggml_tensor * x_prev = ggml_concat(
36 ctx: ctx0, a: att_shift,
37 b: ggml_view_3d(ctx: ctx0, a: att_norm, ne0: n_embd, ne1: n_seq_tokens - 1, ne2: n_seqs, nb1: att_norm->nb[1], nb2: att_norm->nb[2], offset: 0), dim: 1);
38
39 cur = build_rwkv6_time_mix(inp: rs_inp, cur: att_norm, x_prev, ubatch, il);
40
41 ggml_tensor * ffn_inp = ggml_add(ctx: ctx0, a: cur, b: inpL);
42 cb(cur: ffn_inp, name: "ffn_inp", il);
43
44 ggml_tensor * ffn_norm = build_norm(cur: ffn_inp, mw: layer->attn_norm_2, mb: layer->attn_norm_2_b, type: LLM_NORM, il);
45 cb(cur: ffn_norm, name: "ffn_norm", il);
46
47 x_prev = ggml_concat(
48 ctx: ctx0, a: ffn_shift,
49 b: ggml_view_3d(ctx: ctx0, a: ffn_norm, ne0: n_embd, ne1: n_seq_tokens - 1, ne2: n_seqs, nb1: ffn_norm->nb[1], nb2: ffn_norm->nb[2], offset: 0), dim: 1);
50
51 token_shift = ggml_concat(ctx: ctx0,
52 a: ggml_view_3d(ctx: ctx0, a: att_norm, ne0: n_embd, ne1: 1, ne2: n_seqs, nb1: att_norm->nb[1], nb2: att_norm->nb[2],
53 offset: (n_seq_tokens - 1) * n_embd * ggml_element_size(tensor: att_norm)),
54 b: ggml_view_3d(ctx: ctx0, a: ffn_norm, ne0: n_embd, ne1: 1, ne2: n_seqs, nb1: ffn_norm->nb[1], nb2: ffn_norm->nb[2],
55 offset: (n_seq_tokens - 1) * n_embd * ggml_element_size(tensor: ffn_norm)),
56 dim: 1);
57 ggml_build_forward_expand(cgraph: gf, tensor: build_rwkv_token_shift_store(token_shift, ubatch, il));
58
59 ffn_inp = ggml_reshape_2d(ctx: ctx0, a: ffn_inp, ne0: n_embd, ne1: n_tokens);
60 ffn_norm = ggml_reshape_2d(ctx: ctx0, a: ffn_norm, ne0: n_embd, ne1: n_tokens);
61 x_prev = ggml_reshape_2d(ctx: ctx0, a: x_prev, ne0: n_embd, ne1: n_tokens);
62 cur = ggml_reshape_2d(ctx: ctx0, a: cur, ne0: n_embd, ne1: n_tokens);
63
64 if (il == n_layer - 1 && inp_out_ids) {
65 ffn_inp = ggml_get_rows(ctx: ctx0, a: ffn_inp, b: inp_out_ids);
66 ffn_norm = ggml_get_rows(ctx: ctx0, a: ffn_norm, b: inp_out_ids);
67 x_prev = ggml_get_rows(ctx: ctx0, a: x_prev, b: inp_out_ids);
68 cur = ggml_get_rows(ctx: ctx0, a: cur, b: inp_out_ids);
69 }
70 cur = build_rwkv6_channel_mix(layer, cur: ffn_norm, x_prev, arch: LLM_ARCH_RWKV6);
71 cur = ggml_add(ctx: ctx0, a: cur, b: ffn_inp);
72
73 if (hparams.rescale_every_n_layers != 0 && (il + 1) % hparams.rescale_every_n_layers == 0) {
74 cur = ggml_scale(ctx: ctx0, a: cur, s: 0.5F);
75 }
76 cur = build_cvec(cur, il);
77 cb(cur, name: "l_out", il);
78
79 // input for next layer
80 inpL = cur;
81 }
82 cur = inpL;
83 cur = build_norm(cur, mw: model.output_norm, mb: model.output_norm_b, type: LLM_NORM, il: -1);
84
85 cb(cur, name: "result_norm", il: -1);
86 res->t_embd = cur;
87
88 cur = build_lora_mm(w: model.output, cur);
89
90 cb(cur, name: "result_output", il: -1);
91 res->t_logits = cur;
92
93 ggml_build_forward_expand(cgraph: gf, tensor: cur);
94}
95