1#include "models.h"
2
3llm_build_openelm::llm_build_openelm(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
4 const int64_t n_embd_head = hparams.n_embd_head_v;
5
6 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
7
8 ggml_tensor * cur;
9 ggml_tensor * inpL;
10 inpL = build_inp_embd(tok_embd: model.tok_embd);
11
12 // inp_pos - contains the positions
13 ggml_tensor * inp_pos = build_inp_pos();
14
15 auto * inp_attn = build_attn_inp_kv();
16
17 ggml_tensor * inp_out_ids = build_inp_out_ids();
18
19 for (int il = 0; il < n_layer; ++il) {
20 const int64_t n_head = hparams.n_head(il);
21 const int64_t n_head_kv = hparams.n_head_kv(il);
22 const int64_t n_head_qkv = 2*n_head_kv + n_head;
23
24 cur = inpL;
25 ggml_tensor * residual = cur;
26
27 // norm
28 cur = build_norm(cur: inpL,
29 mw: model.layers[il].attn_norm, NULL,
30 type: LLM_NORM_RMS, il);
31 cb(cur, name: "attn_norm", il);
32
33 // self-attention
34 {
35 cur = build_lora_mm(w: model.layers[il].wqkv, cur);
36 cb(cur, name: "wqkv", il);
37
38 cur = ggml_reshape_3d(ctx: ctx0, a: cur, ne0: n_embd_head_k, ne1: n_head_qkv, ne2: n_tokens);
39
40 ggml_tensor * Qcur = ggml_view_3d(ctx: ctx0, a: cur, ne0: n_embd_head, ne1: n_head, ne2: n_tokens, nb1: cur->nb[1], nb2: cur->nb[2], offset: 0);
41 cb(cur: Qcur, name: "Qcur", il);
42
43 ggml_tensor * Kcur = ggml_view_3d(ctx: ctx0, a: cur, ne0: n_embd_head, ne1: n_head_kv, ne2: n_tokens, nb1: cur->nb[1], nb2: cur->nb[2], offset: cur->nb[1]*n_head);
44 cb(cur: Kcur, name: "Kcur", il);
45
46 ggml_tensor * Vcur = ggml_cont(ctx: ctx0, a: ggml_view_3d(ctx: ctx0, a: cur, ne0: n_embd_head, ne1: n_head_kv, ne2: n_tokens, nb1: cur->nb[1], nb2: cur->nb[2], offset: cur->nb[1]*(n_head+n_head_kv)));
47 cb(cur: Vcur, name: "Vcur", il);
48
49 Qcur = build_norm(cur: Qcur,
50 mw: model.layers[il].attn_q_norm, NULL,
51 type: LLM_NORM_RMS, il);
52 cb(cur: Qcur, name: "Qcur", il);
53
54 Kcur = build_norm(cur: Kcur,
55 mw: model.layers[il].attn_k_norm, NULL,
56 type: LLM_NORM_RMS, il);
57 cb(cur: Kcur, name: "Kcur", il);
58
59 Qcur = ggml_rope_ext(
60 ctx: ctx0, a: Qcur, b: inp_pos, NULL,
61 n_dims: n_rot, mode: rope_type, n_ctx_orig, freq_base, freq_scale,
62 ext_factor, attn_factor, beta_fast, beta_slow
63 );
64
65 Kcur = ggml_rope_ext(
66 ctx: ctx0, a: Kcur, b: inp_pos, NULL,
67 n_dims: n_rot, mode: rope_type, n_ctx_orig, freq_base, freq_scale,
68 ext_factor, attn_factor, beta_fast, beta_slow
69 );
70
71 cb(cur: Qcur, name: "Qcur", il);
72 cb(cur: Kcur, name: "Kcur", il);
73 cb(cur: Qcur, name: "Vcur", il);
74
75 cur = build_attn(inp: inp_attn,
76 wo: model.layers[il].wo, NULL,
77 q_cur: Qcur, k_cur: Kcur, v_cur: Vcur, kq_b: nullptr, sinks: nullptr, v_mla: nullptr, kq_scale: 1.0f/sqrtf(x: float(n_embd_head)), il);
78 }
79 if (il == n_layer - 1 && inp_out_ids) {
80 residual = ggml_get_rows(ctx: ctx0, a: residual, b: inp_out_ids);
81 cur = ggml_get_rows(ctx: ctx0, a: cur, b: inp_out_ids);
82 }
83 ggml_tensor * ffn_inp = ggml_add(ctx: ctx0, a: residual, b: cur);
84 cb(cur: ffn_inp, name: "ffn_inp", il);
85
86 // feed-forward network
87 {
88 cur = build_norm(cur: ffn_inp,
89 mw: model.layers[il].ffn_norm, NULL,
90 type: LLM_NORM_RMS, il);
91 cb(cur, name: "ffn_norm", il);
92
93 cur = build_ffn(cur,
94 up: model.layers[il].ffn_up, NULL, NULL,
95 gate: model.layers[il].ffn_gate, NULL, NULL,
96 down: model.layers[il].ffn_down, NULL, NULL,
97 NULL,
98 type_op: LLM_FFN_SILU, type_gate: LLM_FFN_PAR, il);
99 cb(cur, name: "ffn_out", il);
100 }
101 cur = ggml_add(ctx: ctx0, a: cur, b: ffn_inp);
102
103 cur = build_cvec(cur, il);
104 cb(cur, name: "l_out", il);
105
106 inpL = cur;
107 }
108 cur = inpL;
109
110 // norm
111 cur = build_norm(cur,
112 mw: model.output_norm, NULL,
113 type: LLM_NORM_RMS, il: -1);
114
115 cb(cur, name: "result_norm", il: -1);
116 res->t_embd = cur;
117
118 cur = build_lora_mm(w: model.output, cur);
119
120 cb(cur, name: "result_output", il: -1);
121 res->t_logits = cur;
122
123 ggml_build_forward_expand(cgraph: gf, tensor: cur);
124}
125