Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2#include "llama-memory-recurrent.h"3 4void llama_model_minimax_01::load_arch_hparams(llama_model_loader & ml) {5 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);6 ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale);7 8 // we use n_embd_head_la to set recurrent memory n_embd_s9 hparams.n_embd_head_la = hparams.n_embd_head_k_full;10 11 // Mark recurrent layers (lightning attention layers).12 if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {13 uint32_t full_attn_interval = 8;14 ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);15 for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {16 hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);17 }18 }19 20 switch (hparams.n_layer()) {21 case 80: type = LLM_TYPE_456B; break;22 default: type = LLM_TYPE_UNKNOWN;23 }24}25 26void llama_model_minimax_01::load_arch_tensors(llama_model_loader &) {27 LLAMA_LOAD_LOCALS;28 29 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);30 31 // output32 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);33 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);34 35 // if output is NULL, init from the input tok embed36 if (output == NULL) {37 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);38 }39 40 for (int i = 0; i < n_layer; ++i) {41 auto & layer = layers[i];42 43 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);44 45 if (!hparams.is_recr(i)) {46 create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);47 } else {48 layer.attn_norm_2 = create_tensor(tn(LLM_TENSOR_ATTN_NORM_2, "weight", i), {n_embd_head_k * n_head}, 0);49 layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, 3 * n_embd_head_k * n_head}, 0);50 layer.wg = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0);51 }52 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);53 54 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);55 56 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);57 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff, n_expert}, TENSOR_NOT_REQUIRED);58 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff, n_embd, n_expert}, 0);59 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0);60 }61}62 63std::unique_ptr<llm_graph_context> llama_model_minimax_01::build_arch_graph(const llm_graph_params & params) const {64 return std::make_unique<graph>(*this, params);65}66 67class llm_graph_input_la : public llm_graph_input_i {68public:69 llm_graph_input_la(const llama_hparams & hparams) : hparams(hparams) {}70 71 void set_input(const llama_ubatch * ubatch) override {72 // this operates on assumption that we have an equal ubatch split73 74 const int64_t n_head = hparams.n_head();75 const int32_t n_seqs = ubatch->n_seqs;76 const int32_t n_seqs_unq = ubatch->n_seqs_unq;77 const int32_t n_tokens = ubatch->n_tokens;78 const int32_t n_seq_tokens = ubatch->n_seq_tokens;79 80 std::vector<llama_pos> p0(n_seqs_unq);81 std::fill(p0.begin(), p0.end(), std::numeric_limits<llama_pos>::max());82 83 // get lowest token position in a ubatch for each stream84 for (int i = 0; i < n_tokens; ++i) {85 llama_seq_id seq_id = ubatch->seq_id[i][0];86 int32_t seq_idx = ubatch->seq_idx[seq_id];87 llama_pos pos = ubatch->pos[i];88 if (p0[seq_idx] > pos) {89 p0[seq_idx] = pos;90 }91 }92 93 if (inp_slopes) {94 GGML_ASSERT(ggml_backend_buffer_is_host(inp_slopes->buffer));95 96 float * data = (float *) inp_slopes->data;97 98 float start = powf(2, -powf(2, -(log2f(n_head) - 3)));99 float ratio = start;100 101 for (int h = 0; h < n_head; ++h) {102 data[h] = start * powf(ratio, h);103 }104 }105 106 if (inp_q_decay) {107 GGML_ASSERT(ggml_backend_buffer_is_host(inp_q_decay->buffer));108 109 float * slopes = (float *) inp_slopes->data;110 float * data = (float *) inp_q_decay->data;111 112 for (int s = 0; s < n_seqs; ++s) {113 for (int i = 0; i < n_seq_tokens; ++i) {114 llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0];115 int32_t seq_idx = ubatch->seq_idx[seq_id];116 llama_pos pos = ubatch->pos[s * n_seq_tokens + i];117 int pos_rel = pos - p0[seq_idx];118 119 for (int h = 0; h < n_head; ++h) {120 data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (pos_rel + 1);121 }122 }123 }124 }125 126 if (inp_k_decay) {127 GGML_ASSERT(ggml_backend_buffer_is_host(inp_k_decay->buffer));128 129 float * slopes = (float *) inp_slopes->data;130 float * data = (float *) inp_k_decay->data;131 132 for (int s = 0; s < n_seqs; ++s) {133 for (int i = 0; i < n_seq_tokens; ++i) {134 llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0];135 int32_t seq_idx = ubatch->seq_idx[seq_id];136 llama_pos pos = ubatch->pos[s * n_seq_tokens + i];137 int pos_rel = pos - p0[seq_idx];138 139 for (int h = 0; h < n_head; ++h) {140 data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (n_seq_tokens - pos_rel - 1);141 }142 }143 }144 }145 146 if (inp_diag_decay) {147 GGML_ASSERT(ggml_backend_buffer_is_host(inp_diag_decay->buffer));148 149 float * slopes = (float *) inp_slopes->data;150 float * data = (float *) inp_diag_decay->data;151 152 for (int s = 0; s < n_seqs; ++s) {153 for (int h = 0; h < n_head; ++h) {154 for (int j = 0; j < n_seq_tokens; ++j) {155 llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + j][0];156 int32_t seq_idx = ubatch->seq_idx[seq_id];157 llama_pos pos_j = ubatch->pos[s * n_seq_tokens + j];158 int pos_rel_j = pos_j - p0[seq_idx];159 160 for (int i = 0; i < n_seq_tokens; ++i) {161 llama_pos pos_i = ubatch->pos[s * n_seq_tokens + i];162 int pos_rel_i = pos_i - p0[seq_idx];163 164 int index = pos_rel_j - pos_rel_i;165 float s_index = index >= 0 ? -slopes[h] * index : -INFINITY;166 data[seq_idx * n_head * n_seq_tokens * n_seq_tokens + h * n_seq_tokens * n_seq_tokens + j * n_seq_tokens + i] = s_index;167 }168 }169 }170 }171 }172 }173 174 bool can_reuse(const llm_graph_params & params) override {175 bool res = true;176 177 res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens);178 res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens);179 res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens);180 181 return res;182 }183 184 const llama_hparams hparams;185 186 ggml_tensor * inp_slopes = nullptr; // F32 [n_head]187 ggml_tensor * inp_q_decay = nullptr; // F32 [1, n_head, n_batch]188 ggml_tensor * inp_k_decay = nullptr; // F32 [1, n_head, n_batch]189 ggml_tensor * inp_diag_decay = nullptr; // F32 [n_batch, n_batch, n_head]190};191 192llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {193 const int64_t n_embd_head = hparams.n_embd_head_v();194 195 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());196 // GGML_ASSERT(n_embd_head == n_rot); this is wrong in case of minimax, head_dim = 128, n_rot = 64197 198 const int64_t n_seqs = ubatch.n_seqs;199 const int64_t n_seq_tokens = ubatch.n_seq_tokens;200 201 GGML_ASSERT(n_seqs != 0);202 GGML_ASSERT(ubatch.equal_seqs());203 GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);204 205 ggml_tensor * cur;206 ggml_tensor * inpL;207 208 inpL = build_inp_embd(model.tok_embd);209 210 auto * inp_hybrid = build_inp_mem_hybrid();211 auto * inp_rs = inp_hybrid->get_recr();212 213 ggml_tensor * inp_pos = build_inp_pos();214 ggml_tensor * inp_out_ids = build_inp_out_ids();215 216 llm_graph_input_la * la = nullptr;217 218 auto inp = std::make_unique<llm_graph_input_la>(hparams);219 220 inp->inp_slopes = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_head);221 ggml_set_input(inp->inp_slopes);222 cb(inp->inp_slopes, "slopes", -1);223 224 inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs);225 ggml_set_input(inp->inp_q_decay);226 cb(inp->inp_q_decay, "q_decay_exp", -1);227 228 inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs);229 ggml_set_input(inp->inp_k_decay);230 cb(inp->inp_k_decay, "k_decay_exp", -1);231 232 // [TAG_RESERVE_DIAG_DECAY]233 inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs);234 ggml_set_input(inp->inp_diag_decay);235 cb(inp->inp_diag_decay, "diag_decay_exp", -1);236 237 la = (llm_graph_input_la *) res->add_input(std::move(inp));238 239 ggml_tensor * slopes = la->inp_slopes;240 241 for (int il = 0; il < n_layer; ++il) {242 res->t_layer_inp[il] = inpL;243 244 ggml_tensor * inpSA = inpL;245 246 cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);247 cb(cur, "attn_norm", il);248 249 ggml_tensor * residual = cur;250 251 // self_attention252 if (!hparams.is_recr(il)) {253 // softmax attention layer254 255 auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,256 n_embd_head, n_head, n_head_kv, il);257 258 Qcur = ggml_rope_ext(259 ctx0, Qcur, inp_pos, nullptr,260 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,261 ext_factor, attn_factor, beta_fast, beta_slow262 );263 264 Kcur = ggml_rope_ext(265 ctx0, Kcur, inp_pos, nullptr,266 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,267 ext_factor, attn_factor, beta_fast, beta_slow268 );269 270 cb(Qcur, "Qcur", il);271 cb(Kcur, "Kcur", il);272 cb(Vcur, "Vcur", il);273 274 cur = build_attn(inp_hybrid->get_attn(),275 model.layers[il].wo, NULL, model.layers[il].wo_s,276 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);277 } else {278 // lightning attention layer279 280 const auto * mctx_cur = inp_rs->mctx;281 const auto kv_head = mctx_cur->get_head();282 283 // TODO unneeded - any way to make conv states optional in recurrent memory?284 ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);285 ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);286 ggml_build_forward_expand(gf, conv_state_all);287 288 float slope_scale = 1.0 - 1.0 * il / (n_layer - 1) + 1e-5;289 ggml_tensor * slope_rate = ggml_scale(ctx0, slopes, slope_scale);290 cb(slope_rate, "slope_rate", il);291 292 cur = ggml_reshape_4d(ctx0, cur, cur->ne[0], n_seq_tokens, 1, n_seqs);293 294 ggml_tensor * QKVcur = build_lora_mm(model.layers[il].wqkv, cur);295 cb(QKVcur, "QKVcur", il);296 297 QKVcur = ggml_silu(ctx0, QKVcur);298 cb(QKVcur, "QKVcur_silu", il);299 300 QKVcur = ggml_reshape_4d(ctx0, QKVcur, n_embd_head * 3, n_head, n_seq_tokens, n_seqs);301 302 ggml_tensor * Qcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 0*ggml_element_size(QKVcur)*n_embd_head);303 ggml_tensor * Kcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 1*ggml_element_size(QKVcur)*n_embd_head);304 ggml_tensor * Vcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 2*ggml_element_size(QKVcur)*n_embd_head);305 306 cb(Qcur, "Qcur", il);307 cb(Kcur, "Kcur", il);308 cb(Vcur, "Vcur", il);309 310 // get previous KV311 ggml_tensor * la_states_all = mctx_cur->get_s_l(il);312 ggml_tensor * state = build_rs(inp_rs, la_states_all, hparams.n_embd_s(), n_seqs);313 314 ggml_tensor * kv_old = ggml_reshape_4d(ctx0, state, n_embd_head, n_embd_head, n_head, n_seqs);315 cb(kv_old, "kv_old", il);316 317 ggml_tensor * qkv = nullptr;318 ggml_tensor * kv_new = nullptr;319 {320 // lightning attention321 322 ggml_tensor * q_decay_exp = la->inp_q_decay;323 ggml_tensor * k_decay_exp = la->inp_k_decay;324 ggml_tensor * diag_decay_exp = la->inp_diag_decay;325 326 ggml_tensor * q_decay = ggml_exp(ctx0, ggml_scale(ctx0, q_decay_exp, slope_scale));327 cb(q_decay, "q_decay", il);328 ggml_tensor * k_decay = ggml_exp(ctx0, ggml_scale(ctx0, k_decay_exp, slope_scale));329 cb(k_decay, "k_decay", il);330 ggml_tensor * diag_decay = ggml_exp(ctx0, ggml_scale(ctx0, diag_decay_exp, slope_scale));331 cb(diag_decay, "diag_decay", il);332 333 ggml_tensor * q_s = ggml_mul(ctx0, Qcur, q_decay);334 cb(q_s, "q_s", il);335 336 ggml_tensor * q_s_trans = ggml_permute(ctx0, q_s, 0, 2, 1, 3);337 cb(q_s_trans, "q_s_trans", il);338 339 ggml_tensor * qkv_none_diag = ggml_mul_mat(ctx0, kv_old, q_s_trans);340 cb(qkv_none_diag, "qkv_none_diag", il);341 342 ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);343 cb(q_trans, "q_trans", il);344 345 ggml_tensor * k_trans = ggml_permute(ctx0, Kcur, 0, 2, 1, 3);346 cb(k_trans, "k_trans", il);347 348 ggml_tensor * qk = ggml_mul_mat(ctx0, k_trans, q_trans);349 cb(qk, "qk", il);350 351 qk = ggml_mul(ctx0, qk, diag_decay);352 cb(qk, "qk_s", il);353 354 ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3));355 cb(v_trans, "v_trans", il);356 357 ggml_tensor * qkv_diag = ggml_mul_mat(ctx0, v_trans, qk);358 cb(qkv_diag, "qkv_diag", il);359 360 qkv = ggml_add(ctx0, qkv_none_diag, qkv_diag);361 cb(qkv, "qkv", il);362 363 ggml_build_forward_expand(gf, qkv);364 365 ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0*n_seq_tokens);366 cb(slopes_neg, "slopes_neg", il);367 368 ggml_tensor * block_decay = ggml_exp(ctx0, slopes_neg);369 cb(block_decay, "block_decay", il);370 371 ggml_tensor * block_decay_3d = ggml_reshape_3d(ctx0, block_decay, 1, 1, n_head);372 cb(block_decay_3d, "block_decay_3d", il);373 374 ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, block_decay_3d);375 cb(kv_old_s, "kv_old_s", il);376 377 ggml_tensor * k_after_decay = ggml_mul(ctx0, Kcur, k_decay);378 cb(k_after_decay, "k_after_decay", il);379 380 ggml_tensor * k_after_decay_trans = ggml_cont(ctx0, ggml_permute(ctx0, k_after_decay, 1, 2, 0, 3));381 cb(k_after_decay_trans, "k_after_decay_trans", il);382 383 ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_after_decay_trans, v_trans);384 cb(kv_cur, "kv_cur", il);385 386 kv_new = ggml_add(ctx0, kv_old_s, kv_cur);387 cb(kv_new, "kv_new", il);388 }389 390 // store new KV391 ggml_build_forward_expand(gf,392 ggml_cpy(ctx0, kv_new,393 ggml_view_1d(ctx0, la_states_all, hparams.n_embd_s() * n_seqs,394 kv_head * hparams.n_embd_s() * ggml_element_size(la_states_all))));395 396 qkv = ggml_cont(ctx0, ggml_permute(ctx0, qkv, 0, 2, 1, 3));397 cb(qkv, "qkv_permuted", il);398 399 qkv = ggml_reshape_4d(ctx0, qkv, qkv->ne[0]*qkv->ne[1], qkv->ne[2], 1, qkv->ne[3]);400 401 // norm402 ggml_tensor * qkv_norm = build_norm(qkv,403 model.layers[il].attn_norm_2, NULL,404 LLM_NORM_RMS, il);405 cb(qkv_norm, "qkv_norm", il);406 407 ggml_tensor * g = build_lora_mm(model.layers[il].wg, cur);408 cb(g, "g", il);409 410 g = ggml_sigmoid(ctx0, g);411 cb(g, "g_sigm", il);412 413 cur = ggml_mul(ctx0, g, qkv_norm);414 415 cur = build_lora_mm(model.layers[il].wo, cur);416 cb(cur, "attn_out", il);417 418 cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens*n_seqs);419 cb(cur, "attn_out", il);420 }421 422 if (il == n_layer - 1 && inp_out_ids) {423 cur = ggml_get_rows(ctx0, cur, inp_out_ids);424 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);425 residual = ggml_get_rows(ctx0, residual, inp_out_ids);426 }427 428 residual = ggml_scale(ctx0, residual, hparams.f_residual_scale);429 cb(residual, "residual_scaled_attn", il);430 431 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual);432 cb(ffn_inp, "ffn_inp", il);433 434 // MoE branch435 cur = build_norm(ffn_inp,436 model.layers[il].ffn_norm, NULL,437 LLM_NORM_RMS, il);438 cb(cur, "ffn_norm", il);439 440 residual = cur;441 442 cur = build_moe_ffn(cur,443 model.layers[il].ffn_gate_inp,444 model.layers[il].ffn_up_exps,445 model.layers[il].ffn_gate_exps,446 model.layers[il].ffn_down_exps,447 model.layers[il].ffn_exp_probs_b,448 n_expert, n_expert_used,449 LLM_FFN_SILU, true,450 hparams.expert_weights_scale,451 LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,452 il);453 cb(cur, "ffn_moe_out", il);454 455 residual = ggml_scale(ctx0, residual, hparams.f_residual_scale);456 cb(residual, "residual_scaled_ffn", il);457 458 cur = ggml_add(ctx0, cur, residual);459 cb(cur, "ffn_out", il);460 461 cur = build_cvec(cur, il);462 cb(cur, "l_out", il);463 464 // input for next layer465 inpL = cur;466 }467 468 cur = inpL;469 470 cur = build_norm(cur,471 model.output_norm, NULL,472 LLM_NORM_RMS, -1);473 474 cb(cur, "result_norm", -1);475 res->t_embd = cur;476 477 // lm_head478 cur = build_lora_mm(model.output, cur, model.output_s);479 480 cb(cur, "result_output", -1);481 res->t_logits = cur;482 483 ggml_build_forward_expand(gf, cur);484}485