1#include "models.h"
2
3template <bool iswa>
4llm_build_smallthinker<iswa>::llm_build_smallthinker(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params){
5 const int64_t n_embd_head = hparams.n_embd_head_v;
6
7 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k);
8 GGML_ASSERT(n_embd_head == hparams.n_rot);
9
10 ggml_tensor * cur;
11 ggml_tensor * inpL;
12
13 inpL = build_inp_embd(tok_embd: model.tok_embd);
14
15 // inp_pos - contains the positions
16 ggml_tensor * inp_pos = build_inp_pos();
17
18 using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>;
19 inp_attn_type * inp_attn = nullptr;
20
21 if constexpr (iswa) {
22 inp_attn = build_attn_inp_kv_iswa();
23 } else {
24 inp_attn = build_attn_inp_kv();
25 }
26 ggml_tensor * inp_out_ids = build_inp_out_ids();
27
28 for (int il = 0; il < n_layer; ++il) {
29 ggml_tensor * inpSA = inpL;
30 ggml_tensor * probs = nullptr;
31
32 probs = build_lora_mm(w: model.layers[il].ffn_gate_inp, cur: inpL); // [n_expert, n_tokens]
33 cb(cur: probs, name: "ffn_moe_logits", il);
34
35 // norm
36 cur = build_norm(cur: inpL,mw: model.layers[il].attn_norm, NULL, type: LLM_NORM_RMS, il);
37 cb(cur, name: "attn_norm", il);
38
39 // self_attention
40 {
41 // compute Q and K and RoPE them
42 struct ggml_tensor * Qcur = build_lora_mm(w: model.layers[il].wq, cur);
43 cb(cur: Qcur, name: "Qcur", il);
44
45 struct ggml_tensor * Kcur = build_lora_mm(w: model.layers[il].wk, cur);
46 cb(cur: Kcur, name: "Kcur", il);
47
48 struct ggml_tensor * Vcur = build_lora_mm(w: model.layers[il].wv, cur);
49 cb(cur: Vcur, name: "Vcur", il);
50
51 Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
52 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
53 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
54
55 if (hparams.n_no_rope_layer_step == n_layer || il % hparams.n_no_rope_layer_step != 0) {
56 Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
57 ext_factor, attn_factor, beta_fast, beta_slow);
58
59 Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
60 ext_factor, attn_factor, beta_fast, beta_slow);
61 }
62 cb(cur: Qcur, name: "Qcur", il);
63 cb(cur: Kcur, name: "Kcur", il);
64
65 cur = build_attn(inp_attn,
66 model.layers[il].wo, model.layers[il].bo,
67 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f / sqrtf(x: float(n_embd_head)), il);
68 }
69 if (il == n_layer - 1 && inp_out_ids) {
70 cur = ggml_get_rows(ctx0, cur, inp_out_ids);
71 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
72 probs = ggml_get_rows(ctx0, probs, inp_out_ids);
73 }
74 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
75 cb(cur: ffn_inp, name: "ffn_inp", il);
76
77 // MoE branch
78 cur = build_norm(cur: ffn_inp, mw: model.layers[il].ffn_norm, NULL, type: LLM_NORM_RMS, il);
79 cb(cur, name: "ffn_norm", il);
80
81 ggml_tensor * ffn_out =
82 build_moe_ffn(cur,
83 nullptr,
84 model.layers[il].ffn_up_exps,
85 model.layers[il].ffn_gate_exps,
86 model.layers[il].ffn_down_exps,
87 nullptr,
88 n_expert, n_expert_used,
89 LLM_FFN_RELU, true,
90 false, 0.0,
91 static_cast<llama_expert_gating_func_type>(hparams.expert_gating_func),
92 il, probs);
93
94 cb(cur: ffn_out, name: "ffn_out", il);
95 cur = ffn_out;
96
97 cur = ggml_add(ctx0, cur, ffn_inp);
98 cur = build_cvec(cur, il);
99 cb(cur, name: "l_out", il);
100
101 // input for next layer
102 inpL = cur;
103 }
104 cur = inpL;
105
106 cur = build_norm(cur, mw: model.output_norm, NULL, type: LLM_NORM_RMS, il: -1);
107 cb(cur, name: "result_norm", il: -1);
108 res->t_embd = cur;
109
110 // lm_head
111 cur = build_lora_mm(w: model.output, cur);
112 cb(cur, name: "result_output", il: -1);
113 res->t_logits = cur;
114
115 ggml_build_forward_expand(gf, cur);
116}
117
118// Explicit template instantiations
119template struct llm_build_smallthinker<false>;
120template struct llm_build_smallthinker<true>;
121