Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3void llama_model_gemma3n::load_arch_hparams(llama_model_loader & ml) {4 hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;5 load_swa_pattern(ml, 5);6 7 hparams.n_layer_kv_from_start = 20;8 hparams.f_attention_scale = 1.0f;9 10 ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);11 ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);12 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);13 14 switch (hparams.n_layer()) {15 case 30: type = LLM_TYPE_E2B; break;16 case 35: type = LLM_TYPE_E4B; break;17 default: type = LLM_TYPE_UNKNOWN;18 }19}20 21void llama_model_gemma3n::load_arch_tensors(llama_model_loader &) {22 LLAMA_LOAD_LOCALS;23 24 const int64_t n_altup = hparams.n_altup;25 const int64_t laurel_rank = hparams.laurel_rank;26 const int64_t n_embd_altup = hparams.n_embd_altup;27 28 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);29 // if output is NULL, init from the input tok embed30 if (output == NULL) {31 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);32 }33 34 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);35 36 altup_proj = create_tensor(tn(LLM_TENSOR_ALTUP_PROJ, "weight"), {n_embd, n_embd, n_altup - 1}, 0);37 altup_unembd_proj = create_tensor(tn(LLM_TENSOR_ALTUP_UNEMBD_PROJ, "weight"), {n_embd, n_embd, n_altup - 1}, 0);38 39 per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_altup * n_layer, n_vocab}, 0);40 per_layer_model_proj = create_tensor(tn(LLM_TENSOR_PER_LAYER_MODEL_PROJ, "weight", 0), {n_embd, n_embd_altup * n_layer}, 0);41 per_layer_proj_norm = create_tensor(tn(LLM_TENSOR_PER_LAYER_PROJ_NORM, "weight", 0), {n_embd_altup}, 0);42 43 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);44 45 for (int i = 0; i < n_layer; ++i) {46 auto & layer = layers[i];47 48 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);49 50 create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);51 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);52 53 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);54 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);55 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);56 57 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);58 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);59 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);60 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);61 layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);62 63 // altup & laurel64 layer.per_layer_inp_gate = create_tensor(tn(LLM_TENSOR_PER_LAYER_INP_GATE, "weight", i), {n_embd, n_embd_altup}, 0);65 layer.per_layer_proj = create_tensor(tn(LLM_TENSOR_PER_LAYER_PROJ, "weight", i), {n_embd_altup, n_embd}, 0);66 layer.per_layer_post_norm = create_tensor(tn(LLM_TENSOR_PER_LAYER_POST_NORM, "weight", i), {n_embd}, 0);67 layer.altup_correct_coef = create_tensor(tn(LLM_TENSOR_ALTUP_CORRECT_COEF, "weight", i), {n_altup, n_altup}, 0);68 layer.altup_correct_scale = create_tensor(tn(LLM_TENSOR_ALTUP_CORRECT_SCALE, "weight", i), {n_embd}, 0);69 layer.altup_predict_coef = create_tensor(tn(LLM_TENSOR_ALTUP_PREDICT_COEF, "weight", i), {n_altup, n_altup * n_altup}, 0);70 layer.altup_router = create_tensor(tn(LLM_TENSOR_ALTUP_ROUTER, "weight", i), {n_embd, n_altup}, 0);71 layer.altup_router_norm = create_tensor(tn(LLM_TENSOR_ALTUP_ROUTER_NORM, "weight", i), {n_embd}, 0);72 layer.laurel_l = create_tensor(tn(LLM_TENSOR_LAUREL_L, "weight", i), {n_embd, laurel_rank}, 0);73 layer.laurel_r = create_tensor(tn(LLM_TENSOR_LAUREL_R, "weight", i), {laurel_rank, n_embd}, 0);74 layer.laurel_post_norm = create_tensor(tn(LLM_TENSOR_LAUREL_POST_NORM, "weight", i), {n_embd}, 0);75 }76}77 78std::unique_ptr<llm_graph_context> llama_model_gemma3n::build_arch_graph(const llm_graph_params & params) const {79 return std::make_unique<graph>(*this, params);80}81 82// get 2D slice view from a 3D tensor, the idx corresponds to the 3rd dim83static ggml_tensor * gemma3n_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {84 GGML_ASSERT(idx < (int) x->ne[2]);85 return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], ggml_row_size(x->type, x->ne[0]),86 idx * x->ne[0] * x->ne[1] * ggml_element_size(x));87}88 89llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_params & params) :90 llm_graph_context(params),91 model(model),92 n_embd_head(model.hparams.n_embd_head_k()),93 n_embd_altup(model.hparams.n_embd_altup),94 n_altup(model.hparams.n_altup),95 i_altup_act(model.hparams.i_altup_act) {96 ggml_tensor * cur;97 ggml_tensor * inpL;98 99 inpL = build_inp_embd(model.tok_embd);100 101 // important: do not normalize weights for raw embeddings input (i.e. encoded image embeddings)102 inpL = ggml_scale(ctx0, inpL, ubatch.token ? sqrtf(n_embd) : 1.0f);103 cb(inpL, "inp_scaled", -1);104 105 // inp_pos - contains the positions106 ggml_tensor * inp_pos = build_inp_pos();107 108 // TODO: is causal == true correct? might need some changes109 auto * inp_attn = build_attn_inp_kv_iswa();110 111 ggml_tensor * inp_per_layer = build_inp_per_layer();112 ggml_build_forward_expand(gf, inp_per_layer);113 114 // inp_per_layer now has shape: [n_embd_altup, n_tokens, n_layer]115 inp_per_layer = project_per_layer_inputs(inpL, inp_per_layer);116 117 // inpL now has only 1 altup, project it to the rest of the altups118 // these "added" altups will be concat to the last dim of inpL119 {120 ggml_tensor * target_magnitude = calc_magnitude(inpL);121 ggml_tensor * inp_repeated = ggml_repeat_4d(ctx0, inpL, n_embd, n_tokens, n_altup - 1, 1);122 ggml_tensor * altup_added =123 ggml_mul_mat(ctx0, model.altup_proj, inp_repeated); // shape: [n_embd, n_tokens, n_altup - 1]124 ggml_tensor * new_magnitude = calc_magnitude(altup_added);125 altup_added = ggml_div(ctx0, ggml_mul(ctx0, altup_added, target_magnitude), new_magnitude);126 inpL = ggml_concat(ctx0, inpL, altup_added, 2); // shape: [n_embd, n_tokens, n_altup]127 cb(inpL, "inp_stacked", -1);128 }129 // inpL now has shape: [n_embd, n_tokens, n_altup]130 131 for (int il = 0; il < n_layer; ++il) {132 // this block is made to be closely resemble Gemma3p5DecoderLayer on python code133 const float freq_base_l = model.get_rope_freq_base(cparams, il);134 const float freq_scale_l = model.get_rope_freq_scale(cparams, il);135 136 ggml_tensor * cur = inpL; // [n_embd, n_tokens, n_altup]137 ggml_tensor * predictions = altup_predict(cur, il); // [n_embd, n_tokens, n_altup]138 139 // predicted value will go through self-attention and laurel140 ggml_tensor * active_prediction = gemma3n_view_2d_slice(ctx0, predictions, i_altup_act); // [n_embd, n_tokens]141 cur = active_prediction;142 cb(cur, "active_prediction", il);143 144 // norm145 cur = build_norm(cur, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);146 cb(cur, "attn_norm", il);147 148 // laurel149 ggml_tensor * laurel_out = laurel(cur, il); // [n_embd, n_tokens]150 151 // self-attention152 if (hparams.has_kv(il)) {153 auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, n_head, n_head_kv, il);154 155 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);156 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);157 Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps);158 159 cb(Qcur, "Qcur_normed", il);160 cb(Kcur, "Kcur_normed", il);161 cb(Vcur, "Vcur_normed", il);162 163 Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,164 ext_factor, attn_factor, beta_fast, beta_slow);165 166 Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,167 ext_factor, attn_factor, beta_fast, beta_slow);168 169 cb(Qcur, "Qcur_pos", il);170 cb(Kcur, "Kcur_pos", il);171 172 cur = build_attn(inp_attn, model.layers[il].wo,173 NULL, model.layers[il].wo_s, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,174 hparams.f_attention_scale, il);175 } else {176 // reuse KV cache of earlier layers177 ggml_tensor * Qcur;178 if (model.layers[il].wqkv) {179 ggml_tensor * qkv = build_lora_mm(model.layers[il].wqkv, cur);180 const int64_t q_dim = n_embd_head * n_head;181 Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, q_dim, n_tokens, qkv->nb[1], 0));182 } else {183 Qcur = build_lora_mm(model.layers[il].wq, cur);184 }185 cb(Qcur, "Qcur", il);186 Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);187 188 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);189 cb(Qcur, "Qcur_normed", il);190 191 Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,192 ext_factor, attn_factor, beta_fast, beta_slow);193 cb(Qcur, "Qcur_pos", il);194 195 cur = build_attn(inp_attn,196 model.layers[il].wo, NULL, model.layers[il].wo_s,197 Qcur, nullptr, nullptr, nullptr, nullptr, nullptr, hparams.f_attention_scale, il);198 }199 cur = build_norm(cur, model.layers[il].attn_post_norm, NULL, LLM_NORM_RMS, il);200 cb(cur, "attn_post_norm", il);201 202 cur = ggml_add(ctx0, cur, active_prediction); // [n_embd, n_tokens]203 cb(cur, "attn_gated", il);204 205 ggml_tensor * attn_laurel = ggml_scale(ctx0, ggml_add(ctx0, cur, laurel_out),206 1.0f / sqrtf(2.0f)); // [n_embd, n_tokens]207 cb(attn_laurel, "attn_laurel", il);208 209 cur = build_norm(attn_laurel, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);210 cb(cur, "ffn_norm", il);211 212 // feed-forward network213 {214 ggml_tensor * up_proj = build_lora_mm(model.layers[il].ffn_up, cur);215 ggml_tensor * gate_proj = build_lora_mm(model.layers[il].ffn_gate, cur);216 217 if (il < n_layer_sparsity) {218 // apply activation sparsity219 gate_proj = gaussian_topk(gate_proj);220 }221 gate_proj = ggml_gelu(ctx0, gate_proj);222 223 cur = ggml_mul(ctx0, up_proj, gate_proj);224 cur = build_lora_mm(model.layers[il].ffn_down, cur);225 cb(cur, "ffn_out", il);226 }227 cur = build_norm(cur, model.layers[il].ffn_post_norm, NULL, LLM_NORM_RMS, -1);228 cb(cur, "ffn_post_norm", il);229 230 ggml_tensor * attn_ffw_laurel_gated = ggml_add(ctx0, cur, attn_laurel); // [n_embd, n_tokens]231 cb(attn_ffw_laurel_gated, "attn_ffw_laurel_gated", il);232 233 ggml_tensor * corrected = altup_correct(predictions, attn_ffw_laurel_gated, il); // [n_embd, n_tokens, n_altup]234 235 ggml_tensor * first_prediction; // [n_embd, n_tokens]236 {237 first_prediction = gemma3n_view_2d_slice(ctx0, corrected, i_altup_act); // [n_embd, n_tokens]238 first_prediction = ggml_mul(ctx0, first_prediction, model.layers[il].altup_correct_scale);239 first_prediction = build_lora_mm(model.layers[il].per_layer_inp_gate, first_prediction);240 first_prediction = ggml_gelu(ctx0, first_prediction); // [n_embd_altup, n_tokens]241 cb(first_prediction, "first_prediction_gated", il);242 243 ggml_tensor * inp_this_layer = gemma3n_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_altup, n_tokens]244 first_prediction = ggml_mul(ctx0, first_prediction, inp_this_layer); // [n_embd_altup, n_tokens]245 cb(first_prediction, "first_prediction_scaled", il);246 247 first_prediction = build_lora_mm(model.layers[il].per_layer_proj, first_prediction); // [n_embd, n_tokens]248 first_prediction =249 build_norm(first_prediction, model.layers[il].per_layer_post_norm, NULL, LLM_NORM_RMS, il);250 cb(first_prediction, "first_prediction_out", il);251 }252 // equivalent to python code: corrected_predictions[1:] += first_prediction253 {254 ggml_tensor * slice_first = gemma3n_view_2d_slice(ctx0, corrected, 0);255 ggml_tensor * slice_rest = ggml_view_3d(256 ctx0, corrected, n_embd, n_tokens, n_altup - 1, ggml_row_size(corrected->type, n_embd),257 ggml_row_size(corrected->type, n_embd * n_tokens), n_embd * n_tokens * ggml_element_size(corrected));258 ggml_tensor * tmp = ggml_add(ctx0, slice_rest, first_prediction); // [n_embd, n_tokens, n_altup - 1]259 corrected = ggml_concat(ctx0, slice_first, tmp, 2); // [n_embd, n_tokens, n_altup]260 }261 cur = corrected; // [n_embd, n_tokens, n_altup]262 cur = build_cvec(cur, il);263 cb(cur, "l_out", il);264 265 // input for next layer266 inpL = cur;267 }268 cur = inpL; // [n_embd, n_tokens, n_altup]269 270 // cur now has multiple altup(s), we want to merge them back to 1 altup271 {272 ggml_tensor * target_magnitude = calc_magnitude(gemma3n_view_2d_slice(ctx0, cur, i_altup_act)); // [n_embd, n_tokens]273 // do a view to skip the first slice (active altup)274 ggml_tensor * alt_slice =275 ggml_view_3d(ctx0, cur, n_embd, n_tokens, n_altup - 1, ggml_row_size(cur->type, n_embd),276 ggml_row_size(cur->type, n_embd * n_tokens), n_embd * n_tokens * ggml_element_size(cur));277 ggml_tensor * altup_unembd =278 ggml_mul_mat(ctx0, model.altup_unembd_proj, alt_slice); // shape: [n_embd, n_tokens, n_altup - 1]279 ggml_tensor * new_magnitude = calc_magnitude(altup_unembd);280 altup_unembd = ggml_div(ctx0, ggml_mul(ctx0, altup_unembd, target_magnitude), new_magnitude);281 cb(altup_unembd, "altup_unembd", -1);282 283 // equivalent to torch.mean(hidden_states, dim=0)284 cur = gemma3n_view_2d_slice(ctx0, cur, 0); // [n_embd, n_tokens]285 for (int i = 0; i < n_altup - 1; ++i) {286 cur = ggml_add(ctx0, cur, gemma3n_view_2d_slice(ctx0, altup_unembd, i));287 }288 cur = ggml_scale(ctx0, cur, 1.0f / float(n_altup)); // [n_embd, n_tokens]289 cb(cur, "unembd_merged", -1);290 }291 // cur now has shape: [n_embd, n_tokens]292 293 // TODO: move this to right after the last KV layer294 {295 // skip computing output for unused tokens296 ggml_tensor * inp_out_ids = build_inp_out_ids();297 cur = ggml_get_rows(ctx0, cur, inp_out_ids);298 }299 cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);300 301 cb(cur, "result_norm", -1);302 res->t_embd = cur;303 304 cur = build_lora_mm(model.output, cur, model.output_s);305 306 {307 // final logit soft-capping308 cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);309 cur = ggml_tanh(ctx0, cur);310 cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);311 }312 cb(cur, "result_output", -1);313 res->t_logits = cur;314 315 ggml_build_forward_expand(gf, cur);316}317 318ggml_tensor * llama_model_gemma3n::graph::calc_magnitude(ggml_tensor * x) {319 return ggml_sqrt(ctx0, ggml_sum_rows(ctx0, ggml_sqr(ctx0, x)));320}321 322// equivalent to get_per_layer_inputs() in python code323// output shape: [n_embd_altup, n_layer, n_tokens]324ggml_tensor * llama_model_gemma3n::graph::build_inp_per_layer() {325 auto inp = std::make_unique<llm_graph_input_embd>(n_embd);326 ggml_tensor * inp_per_layer;327 float tok_embd_scale = sqrtf((float) n_embd_altup);328 if (ubatch.token) {329 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens);330 ggml_set_input(inp->tokens);331 res->t_inp_tokens = inp->tokens;332 inp_per_layer = ggml_get_rows (ctx0, model.per_layer_tok_embd, inp->tokens);333 inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_altup, n_layer, n_tokens);334 inp_per_layer = ggml_scale (ctx0, inp_per_layer, tok_embd_scale);335 cb(inp_per_layer, "inp_per_layer_selected", -1);336 res->add_input(std::move(inp));337 } else {338 // Multimodal embedding path: use padding token (ID=0) embedding339 // TODO: verify if this is the correct behavior in transformers implementation340 const int64_t embd_size = model.per_layer_tok_embd->ne[0]; // n_embd_altup * n_layer341 342 // Extract and dequantize padding token embedding (row 0)343 ggml_tensor * padding = ggml_view_1d(ctx0, model.per_layer_tok_embd, embd_size, 0);344 inp_per_layer = ggml_cast (ctx0, padding, GGML_TYPE_F32);345 inp_per_layer = ggml_scale(ctx0, inp_per_layer, tok_embd_scale);346 347 // Reshape to [n_embd_altup, n_layer, 1]348 inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_altup, n_layer, 1);349 cb(inp_per_layer, "inp_per_layer_multimodal", -1);350 }351 return inp_per_layer;352}353 354// equivalent to project_per_layer_inputs() in python code355// this calculates the per-layer inputs, so the final tensor shape will have n_layer as the last dim356// output shape: [n_embd_altup, n_tokens, n_layer]357ggml_tensor * llama_model_gemma3n::graph::project_per_layer_inputs(ggml_tensor * inp_batch, ggml_tensor * inp_per_layer) {358 const float per_layer_projection_scale = 1.0f / sqrtf((float) n_embd);359 const float per_layer_input_scale = 1.0f / sqrtf(2.0f);360 361 ggml_tensor * per_layer_proj;362 per_layer_proj = ggml_mul_mat (ctx0, model.per_layer_model_proj, inp_batch);363 per_layer_proj = ggml_scale (ctx0, per_layer_proj, per_layer_projection_scale);364 per_layer_proj = ggml_reshape_3d(ctx0, per_layer_proj, n_embd_altup, n_layer, n_tokens);365 366 per_layer_proj = build_norm(per_layer_proj, model.per_layer_proj_norm, NULL, LLM_NORM_RMS, -1);367 cb(per_layer_proj, "per_layer_proj", -1);368 369 inp_per_layer = ggml_add (ctx0, per_layer_proj, inp_per_layer);370 inp_per_layer = ggml_scale(ctx0, inp_per_layer, per_layer_input_scale);371 cb(inp_per_layer, "inp_per_layer", -1);372 373 // permute to shape: [n_embd_altup, n_tokens, n_layer]374 inp_per_layer = ggml_cont(ctx0, ggml_permute(ctx0, inp_per_layer, 0, 2, 1, 3));375 return inp_per_layer;376}377 378// input cur shape: [n_altup, n_tokens]379// output shape: [n_altup, n_tokens]380ggml_tensor * llama_model_gemma3n::graph::laurel(ggml_tensor * cur, int il) {381 ggml_tensor * tmp = cur;382 tmp = build_lora_mm(model.layers[il].laurel_l, tmp);383 tmp = build_lora_mm(model.layers[il].laurel_r, tmp);384 tmp = build_norm(tmp, model.layers[il].laurel_post_norm, NULL, LLM_NORM_RMS, il);385 tmp = ggml_add(ctx0, tmp, cur);386 cb(tmp, "laurel_out", il);387 return tmp;388}389 390// input x shape: [n_embd, n_tokens]391// output shape: [n_embd, n_tokens]392ggml_tensor * llama_model_gemma3n::graph::gaussian_topk(ggml_tensor * x) {393 ggml_tensor * mean = ggml_mean(ctx0, x);394 ggml_tensor * std = ggml_sqrt(ctx0, ggml_scale(ctx0, ggml_sum_rows(ctx0, ggml_sqr(ctx0, ggml_sub(ctx0, x, mean))),395 1.0f / (float) (x->ne[0] - 1)));396 ggml_tensor * cutoff_x = ggml_add(ctx0, mean, ggml_scale(ctx0, std, f_sparsity_std_mul));397 return ggml_relu(ctx0, ggml_sub(ctx0, x, cutoff_x));398}399 400//401// altup functions402//403 404// equivalent to compute_router_modalities() in python code405// input x shape: [n_embd, n_tokens]406// output shape: [n_altup, n_tokens]407ggml_tensor * llama_model_gemma3n::graph::altup_compute_router_modalities(ggml_tensor * x, int il) {408 ggml_tensor * router_inputs = build_norm(x, model.layers[il].altup_router_norm, NULL, LLM_NORM_RMS, il);409 410 // router_input_scale411 router_inputs = ggml_scale(ctx0, router_inputs, 1.0f / (float) n_embd);412 413 ggml_tensor * output = ggml_mul_mat(ctx0, model.layers[il].altup_router, router_inputs);414 return ggml_tanh(ctx0, output); // [n_altup, n_tokens]415}416 417// input cur shape: [n_embd, n_tokens, n_altup]418// output shape: [n_embd, n_tokens, n_altup]419ggml_tensor * llama_model_gemma3n::graph::altup_predict(ggml_tensor * cur, int il) {420 ggml_tensor * activated = gemma3n_view_2d_slice(ctx0, cur, i_altup_act); // [n_embd, n_tokens]421 ggml_tensor * modalities = altup_compute_router_modalities(activated, il); // [n_altup, n_tokens]422 cb(modalities, "modalities", il);423 424 ggml_tensor * all_coefs = build_lora_mm(model.layers[il].altup_predict_coef, modalities);425 cb(all_coefs, "all_coefs", il);426 // first dim now having n_altup^2 elements, we reshape it to 2D (so we end up with 3D tensor)427 all_coefs = ggml_reshape_3d(ctx0, all_coefs, n_altup, n_altup, n_tokens);428 429 // permute to [n_altup, n_embd, n_tokens]430 ggml_tensor * cur_permuted = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3));431 ggml_tensor * predictions = ggml_mul_mat(ctx0, cur_permuted, all_coefs); // [n_altup, n_embd, n_tokens]432 433 // final shape must be the same as cur: [n_embd, n_tokens, n_altup]434 predictions = ggml_cont(ctx0, ggml_permute(ctx0, predictions, 0, 2, 1, 3));435 predictions = ggml_add(ctx0, predictions, cur);436 cb(predictions, "predictions", il);437 438 return predictions;439}440 441// input predictions shape: [n_embd, n_tokens, n_altup]442// input activated shape: [n_embd, n_tokens]443// output shape: [n_embd, n_tokens, n_altup]444ggml_tensor * llama_model_gemma3n::graph::altup_correct(ggml_tensor * predictions, ggml_tensor * activated, int il) {445 ggml_tensor * modalities = altup_compute_router_modalities(activated, il); // [n_altup, n_tokens]446 cb(modalities, "modalities", il);447 448 ggml_tensor * active_prediction = gemma3n_view_2d_slice(ctx0, predictions, i_altup_act);449 ggml_tensor * innovation = ggml_sub(ctx0, activated, active_prediction); // [n_embd, n_tokens]450 cb(innovation, "innovation", il);451 452 ggml_tensor * all_coefs = build_lora_mm(model.layers[il].altup_correct_coef, modalities); // [n_altup, n_tokens]453 all_coefs = ggml_scale_bias(ctx0, all_coefs, 1.0f, 1.0f); // + 1.0454 cb(all_coefs, "all_coefs", il);455 all_coefs = ggml_transpose(ctx0, all_coefs); // [n_tokens, n_altup]456 all_coefs = ggml_cont_3d(ctx0, all_coefs, 1, n_tokens, n_altup); // [1, n_tokens, n_altup]457 458 innovation = ggml_repeat_4d(ctx0, innovation, n_embd, n_tokens, n_altup, 1);459 ggml_tensor * corrected = ggml_mul(ctx0, innovation, all_coefs); // [n_embd, n_tokens, n_altup]460 corrected = ggml_add(ctx0, corrected, predictions); // [n_embd, n_tokens, n_altup]461 cb(corrected, "corrected", il);462 463 return corrected;464}465