Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3void llama_model_gemma4::load_arch_hparams(llama_model_loader & ml) {4 hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;5 ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);6 7 uint32_t n_kv_shared_layers = 0;8 ml.get_key(LLM_KV_ATTENTION_SHARED_KV_LAYERS, n_kv_shared_layers, false);9 10 hparams.n_layer_kv_from_start = hparams.n_layer_all - (int32_t)n_kv_shared_layers;11 hparams.f_attention_scale = 1.0f; // Gemma4 uses self.scaling = 1.0 (no pre-attn scaling)12 13 ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);14 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all, false);15 ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);16 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);17 ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer);18 ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_SWA, hparams.n_embd_head_k_swa);19 ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_SWA, hparams.n_embd_head_v_swa);20 ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);21 22 // when non_causal is set, the model will use bidirectional attention on SWA layers only, while dense layers will remain causal23 // ref: use_bidirectional_attention == "vision" in HF config24 // note: E2B/E4B are always causal, bypassing this logic25 hparams.non_causal_type = LLAMA_NON_CAUSAL_TYPE_SWA_ONLY;26 27 switch (hparams.n_layer()) {28 case 30: type = LLM_TYPE_26B_A4B; break;29 case 35: type = LLM_TYPE_E2B; break;30 case 42: type = LLM_TYPE_E4B; break;31 case 60: type = LLM_TYPE_31B; break;32 default: type = LLM_TYPE_UNKNOWN;33 }34}35 36void llama_model_gemma4::load_arch_tensors(llama_model_loader &) {37 LLAMA_LOAD_LOCALS;38 39 const uint32_t n_embd_per_layer = hparams.n_embd_per_layer;40 const int64_t n_ff_exp = hparams.n_ff_exp();41 42 if (n_embd_head_k != n_embd_head_v) {43 throw std::runtime_error("Gemma 4 requires n_embd_head_k == n_embd_head_v");44 }45 if (hparams.n_embd_head_k_swa != hparams.n_embd_head_v_swa) {46 throw std::runtime_error("Gemma 4 requires n_embd_head_k_swa == n_embd_head_v_swa");47 }48 49 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);50 // if output is NULL, init from the input tok embed51 if (output == NULL) {52 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);53 }54 55 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);56 57 if (n_embd_per_layer > 0) {58 per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, TENSOR_READ_LAZY);59 per_layer_model_proj = create_tensor(tn(LLM_TENSOR_PER_LAYER_MODEL_PROJ, "weight", 0), {n_embd, n_embd_per_layer * n_layer}, 0);60 per_layer_proj_norm = create_tensor(tn(LLM_TENSOR_PER_LAYER_PROJ_NORM, "weight", 0), {n_embd_per_layer}, 0);61 }62 63 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);64 65 int rope_freqs_flag = 0;66 67 for (int i = 0; i < n_layer; ++i) {68 auto & layer = layers[i];69 const int64_t n_head = hparams.n_head(i);70 const int64_t n_embd_head = hparams.n_embd_head_k(i);71 const int64_t n_embd_k = hparams.n_embd_k_gqa(i);72 const int64_t n_embd_v = hparams.n_embd_v_gqa(i);73 const int kv_flags = hparams.has_kv(i) ? 0 : TENSOR_NOT_REQUIRED;74 75 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);76 77 // note: use_alternative_attention (v_proj is optional, if it's not present, use k_proj)78 layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i),79 {n_embd, n_embd_head * n_head + n_embd_k + n_embd_v}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);80 if (!layer.wqkv) {81 layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head * n_head}, 0);82 layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k}, kv_flags);83 layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v}, TENSOR_NOT_REQUIRED);84 }85 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head * n_head, n_embd}, 0);86 87 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head}, 0);88 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head}, kv_flags);89 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0);90 91 layer.out_scale = create_tensor(tn(LLM_TENSOR_LAYER_OUT_SCALE, "weight", i), {1u}, TENSOR_NOT_REQUIRED);92 93 if (!hparams.is_swa(i)) {94 // full_attention layers use rope_freqs for proportional rope95 layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_embd_head/2}, rope_freqs_flag);96 rope_freqs_flag = TENSOR_DUPLICATED;97 }98 99 // handle use_double_wide_mlp100 int64_t n_ff_cur = hparams.n_ff(i);101 102 // for expert layers, we use normal FFN as shared expert (same as python code)103 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);104 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff_cur}, 0);105 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff_cur}, 0);106 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff_cur, n_embd}, 0);107 layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0);108 109 // MoE router110 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED);111 bool has_expert = layer.ffn_gate_inp != nullptr;112 113 // norm114 if (has_expert) {115 layer.ffn_gate_inp_s = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "scale", i), {n_embd}, 0);116 117 layer.ffn_pre_norm_2 = create_tensor(tn(LLM_TENSOR_FFN_PRE_NORM_2, "weight", i), {n_embd}, 0);118 layer.ffn_post_norm_1 = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM_1, "weight", i), {n_embd}, 0);119 layer.ffn_post_norm_2 = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM_2, "weight", i), {n_embd}, 0);120 121 // MoE FFN122 layer.ffn_gate_up_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_UP_EXPS, "weight", i), {n_embd, n_ff_exp * 2, n_expert}, TENSOR_NOT_REQUIRED);123 124 if (layer.ffn_gate_up_exps == nullptr) {125 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);126 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);127 }128 129 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);130 131 // per-expert scale will be loaded as down_exps_s at the end of the current switch case132 }133 134 // per-layer embeddings135 if (n_embd_per_layer > 0) {136 layer.per_layer_inp_gate = create_tensor(tn(LLM_TENSOR_PER_LAYER_INP_GATE, "weight", i), {n_embd, n_embd_per_layer}, 0);137 layer.per_layer_proj = create_tensor(tn(LLM_TENSOR_PER_LAYER_PROJ, "weight", i), {n_embd_per_layer, n_embd}, 0);138 layer.per_layer_post_norm = create_tensor(tn(LLM_TENSOR_PER_LAYER_POST_NORM, "weight", i), {n_embd}, 0);139 }140 }141}142 143std::unique_ptr<llm_graph_context> llama_model_gemma4::build_arch_graph(const llm_graph_params & params) const {144 return std::make_unique<graph>(*this, params);145}146 147// get 2D slice view from a 3D tensor, the idx corresponds to the 3rd dim148static ggml_tensor * gemma4_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {149 GGML_ASSERT(idx < (int) x->ne[2]);150 return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], ggml_row_size(x->type, x->ne[0]),151 idx * x->ne[0] * x->ne[1] * ggml_element_size(x));152}153 154llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) :155 llm_graph_context(params),156 model(model),157 n_embd_per_layer(model.hparams.n_embd_per_layer) {158 ggml_tensor * cur;159 ggml_tensor * inpL;160 161 inpL = build_inp_embd(model.tok_embd);162 163 // important: do not normalize weights for raw embeddings input (i.e. encoded image emdeddings)164 inpL = ggml_scale(ctx0, inpL, ubatch.token ? sqrtf(n_embd) : 1.0f);165 cb(inpL, "inp_scaled", -1);166 167 // inp_pos - contains the positions168 ggml_tensor * inp_pos = build_inp_pos();169 170 // TODO: is causal == true correct? might need some changes171 auto * inp_attn = build_attn_inp_kv_iswa();172 173 ggml_tensor * inp_out_ids = build_inp_out_ids();174 175 ggml_tensor * inp_per_layer = nullptr;176 if (model.per_layer_tok_embd) {177 inp_per_layer = build_inp_per_layer();178 ggml_build_forward_expand(gf, inp_per_layer);179 180 // inp_per_layer shape: [n_embd_per_layer, n_tokens, n_layer]181 inp_per_layer = project_per_layer_inputs(inpL, inp_per_layer);182 }183 184 for (int il = 0; il < n_layer; ++il) {185 const int64_t n_embd_head = hparams.n_embd_head_k(il);186 GGML_ASSERT(n_embd_head == hparams.n_embd_head_v(il));187 188 const int64_t n_head = hparams.n_head(il);189 const int64_t n_head_kv = hparams.n_head_kv(il);190 191 const float freq_base_l = model.get_rope_freq_base(cparams, il);192 const float freq_scale_l = model.get_rope_freq_scale(cparams, il);193 const int n_rot_l = hparams.n_rot(il);194 195 res->t_layer_inp[il] = inpL;196 197 // norm198 cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);199 cb(cur, "attn_norm", il);200 201 ggml_tensor * freq_factors = nullptr;202 if (!hparams.is_swa(il)) {203 // full_attention layers use rope_freqs for proportional rope204 freq_factors = model.layers[il].rope_freqs;205 }206 207 // Q projection (shared for both non-KV and KV layers)208 // this is to mirror Gemma4Attention in pytorch code209 ggml_tensor * qkv_fused = nullptr;210 ggml_tensor * Qcur;211 if (model.layers[il].wqkv) {212 qkv_fused = build_lora_mm(model.layers[il].wqkv, cur, model.layers[il].wqkv_s);213 cb(qkv_fused, "wqkv", il);214 const int64_t q_dim = n_embd_head * n_head;215 Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, q_dim, n_tokens, qkv_fused->nb[1], 0));216 } else {217 Qcur = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s);218 }219 {220 cb(Qcur, "Qcur", il);221 222 Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);223 224 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);225 cb(Qcur, "Qcur_normed", il);226 227 Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, freq_factors, n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,228 ext_factor, attn_factor, beta_fast, beta_slow);229 cb(Qcur, "Qcur_pos", il);230 }231 232 // self-attention233 if (hparams.has_kv(il)) {234 ggml_tensor * Kcur;235 ggml_tensor * Vcur;236 if (qkv_fused) {237 const int64_t q_dim = n_embd_head * n_head;238 const int64_t k_dim = n_embd_head * n_head_kv;239 const int64_t v_dim = n_embd_head * n_head_kv;240 const size_t esize = ggml_element_size(qkv_fused);241 Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, k_dim, n_tokens, qkv_fused->nb[1], q_dim * esize));242 Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, v_dim, n_tokens, qkv_fused->nb[1], (q_dim + k_dim) * esize));243 } else {244 Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s);245 Vcur = model.layers[il].wv246 ? build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s)247 : Kcur; // if v_proj is not present, use Kcur as Vcur248 }249 cb(Kcur, "Kcur", il);250 cb(Vcur, "Vcur", il);251 252 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);253 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);254 255 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);256 Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps);257 258 cb(Kcur, "Kcur_normed", il);259 cb(Vcur, "Vcur_normed", il);260 261 Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, freq_factors, n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,262 ext_factor, attn_factor, beta_fast, beta_slow);263 264 cb(Kcur, "Kcur_pos", il);265 266 cur = build_attn(inp_attn, model.layers[il].wo,267 nullptr, model.layers[il].wo_s, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,268 hparams.f_attention_scale, il);269 } else {270 // reuse KV cache of earlier layers271 cur = build_attn(inp_attn,272 model.layers[il].wo, nullptr, model.layers[il].wo_s,273 Qcur, nullptr, nullptr, nullptr, nullptr, nullptr, hparams.f_attention_scale, il);274 }275 276 // TODO @ngxson : strip unused token right after the last KV layer to speed up prompt processing277 // keep all rows when extracting unmasked nextn embeddings (MTP target needs the hidden state for every token)278 if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {279 cur = ggml_get_rows(ctx0, cur, inp_out_ids);280 inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);281 }282 cur = build_norm(cur,283 model.layers[il].attn_post_norm, nullptr,284 LLM_NORM_RMS, il);285 cb(cur, "attn_post_norm", il);286 287 ggml_tensor * attn_out = ggml_add(ctx0, cur, inpL);288 cb(attn_out, "attn_out", il);289 290 // feed-forward network291 const bool is_moe_layer = model.layers[il].ffn_gate_inp != nullptr;292 if (is_moe_layer) {293 // MLP (shared exp)294 ggml_tensor * cur_mlp = build_norm(attn_out,295 model.layers[il].ffn_norm, nullptr,296 LLM_NORM_RMS, il);297 cb(cur_mlp, "ffn_norm_1", il);298 299 cur_mlp = build_ffn(cur_mlp,300 model.layers[il].ffn_up, nullptr, model.layers[il].ffn_up_s,301 model.layers[il].ffn_gate, nullptr, model.layers[il].ffn_gate_s,302 model.layers[il].ffn_down, nullptr, model.layers[il].ffn_down_s,303 nullptr,304 LLM_FFN_GELU, LLM_FFN_PAR, il);305 cur_mlp = build_norm(cur_mlp,306 model.layers[il].ffn_post_norm_1, nullptr,307 LLM_NORM_RMS, il);308 cb(cur_mlp, "ffn_mlp", il);309 310 // Expert FFN311 ggml_tensor * cur_moe = build_norm(attn_out,312 model.layers[il].ffn_pre_norm_2, nullptr,313 LLM_NORM_RMS, il);314 cb(cur_moe, "ffn_norm_2", il);315 316 // custom MoE logits calculation (router operates on attn_out, not cur)317 ggml_tensor * tmp = ggml_rms_norm(ctx0, attn_out, hparams.f_norm_rms_eps);318 tmp = ggml_scale(ctx0, tmp, 1.0f / sqrtf((float) n_embd));319 tmp = ggml_mul(ctx0, tmp, model.layers[il].ffn_gate_inp_s);320 ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, tmp); // [n_expert, n_tokens]321 cb(logits, "ffn_moe_logits", il);322 323 cur_moe = build_moe_ffn(cur_moe,324 nullptr, // gate_inp325 model.layers[il].ffn_up_exps,326 model.layers[il].ffn_gate_exps,327 model.layers[il].ffn_down_exps,328 nullptr, // exp_probs_b (not used for gemma4)329 n_expert, n_expert_used,330 LLM_FFN_GELU, true,331 1.0f,332 LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,333 il, logits,334 model.layers[il].ffn_gate_up_exps,335 model.layers[il].ffn_up_exps_s,336 model.layers[il].ffn_gate_exps_s,337 model.layers[il].ffn_down_exps_s);338 cur_moe = build_norm(cur_moe,339 model.layers[il].ffn_post_norm_2, nullptr,340 LLM_NORM_RMS, il);341 cb(cur_moe, "ffn_moe", il);342 343 cur = ggml_add(ctx0, cur_mlp, cur_moe);344 cb(cur, "ffn_moe_combined", il);345 } else {346 cur = build_norm(attn_out,347 model.layers[il].ffn_norm, nullptr,348 LLM_NORM_RMS, il);349 cb(cur, "ffn_norm", il);350 351 cur = build_ffn(cur,352 model.layers[il].ffn_up, nullptr, model.layers[il].ffn_up_s,353 model.layers[il].ffn_gate, nullptr, model.layers[il].ffn_gate_s,354 model.layers[il].ffn_down, nullptr, model.layers[il].ffn_down_s,355 nullptr,356 LLM_FFN_GELU, LLM_FFN_PAR, il);357 cb(cur, "ffn_out", il);358 }359 cur = build_norm(cur,360 model.layers[il].ffn_post_norm, nullptr,361 LLM_NORM_RMS, -1);362 cb(cur, "ffn_post_norm", il);363 364 // residual connection365 cur = ggml_add(ctx0, cur, attn_out);366 367 // per-layer embedding368 if (inp_per_layer) {369 ggml_tensor * pe_in = cur;370 cb(cur, "pe_in", il);371 372 cur = build_lora_mm(model.layers[il].per_layer_inp_gate, cur); // [n_embd_per_layer, n_tokens]373 cur = ggml_gelu(ctx0, cur);374 375 ggml_tensor * inp_this_layer = gemma4_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]376 377 // TODO @ngxson : improve this378 if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {379 inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids);380 }381 382 cur = ggml_mul(ctx0, cur, inp_this_layer);383 cur = build_lora_mm(model.layers[il].per_layer_proj, cur); // [n_embd, n_tokens]384 cur = build_norm(cur, model.layers[il].per_layer_post_norm, nullptr, LLM_NORM_RMS, il);385 cb(cur, "per_layer_embd_out", il);386 387 // residual connection388 cur = ggml_add(ctx0, pe_in, cur);389 }390 391 // layer_scalar392 if (model.layers[il].out_scale) {393 cur = ggml_mul(ctx0, cur, model.layers[il].out_scale);394 cb(cur, "out_scaled", il);395 }396 397 cur = build_cvec(cur, il);398 cb(cur, "l_out", il);399 400 // input for next layer401 inpL = cur;402 }403 cur = inpL;404 405 cur = build_norm(cur,406 model.output_norm, nullptr,407 LLM_NORM_RMS, -1);408 409 // Expose the post-output-norm hidden state (the LM-head input feature) so that410 // MTP draft contexts can read it via llama_get_embeddings_nextn_ith() as the411 // recurrent h input. This matches the reference (transformers/vLLM/SGLang),412 // which feeds the drafter the target's post-final-norm hidden state.413 cb(cur, "h_nextn", -1);414 res->t_h_nextn = cur;415 416 if (!cparams.embeddings_nextn_masked && inp_out_ids) {417 cur = ggml_get_rows(ctx0, cur, inp_out_ids);418 }419 420 cb(cur, "result_norm", -1);421 res->t_embd = cur;422 423 // lm_head424 cur = build_lora_mm(model.output, cur, model.output_s);425 426 if (hparams.f_final_logit_softcapping) {427 cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);428 cur = ggml_tanh(ctx0, cur);429 cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);430 }431 432 cb(cur, "result_output", -1);433 res->t_logits = cur;434 435 ggml_build_forward_expand(gf, cur);436}437 438// equivalent to get_per_layer_inputs() in python code439// output shape: [n_embd_per_layer, n_layer, n_tokens]440ggml_tensor * llama_model_gemma4::graph::build_inp_per_layer() {441 auto inp = std::make_unique<llm_graph_input_embd>(n_embd);442 443 ggml_tensor * inp_per_layer;444 float tok_embd_scale = sqrtf((float) n_embd_per_layer);445 if (ubatch.token) {446 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens);447 ggml_set_input(inp->tokens);448 res->t_inp_tokens = inp->tokens;449 450 inp_per_layer = ggml_get_rows (ctx0, model.per_layer_tok_embd, inp->tokens);451 inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_per_layer, n_layer, n_tokens);452 inp_per_layer = ggml_scale (ctx0, inp_per_layer, tok_embd_scale);453 cb(inp_per_layer, "inp_per_layer_selected", -1);454 455 res->add_input(std::move(inp));456 } else {457 // Multimodal embedding path: use padding token (ID=0) embedding458 // TODO: verify if this is the correct behavior in transformers implementation459 const int64_t embd_size = model.per_layer_tok_embd->ne[0]; // n_embd_per_layer * n_layer460 461 // Extract and dequantize padding token embedding (row 0)462 ggml_tensor * padding = ggml_view_1d(ctx0, model.per_layer_tok_embd, embd_size, 0);463 inp_per_layer = ggml_cast (ctx0, padding, GGML_TYPE_F32);464 inp_per_layer = ggml_scale(ctx0, inp_per_layer, tok_embd_scale);465 466 // Reshape to [n_embd_per_layer, n_layer, 1]467 inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_per_layer, n_layer, 1);468 cb(inp_per_layer, "inp_per_layer_multimodal", -1);469 }470 return inp_per_layer;471}472 473// equivalent to project_per_layer_inputs() in python code474// this calculates the per-layer inputs, so the final tensor shape will have n_layer as the last dim475// inp_batch shape: [n_embd, n_tokens]476// inp_per_layer shape: [n_embd_per_layer, n_layer, n_tokens] (from build_inp_per_layer)477// output shape: [n_embd_per_layer, n_tokens, n_layer]478ggml_tensor * llama_model_gemma4::graph::project_per_layer_inputs(ggml_tensor * inp_batch, ggml_tensor * inp_per_layer) {479 const float per_layer_projection_scale = 1.0f / sqrtf((float) n_embd);480 const float per_layer_input_scale = 1.0f / sqrtf(2.0f);481 482 // note: this matrix multiplication will be performed in the input layer (i.e. on the CPU)483 ggml_tensor * per_layer_proj;484 per_layer_proj = ggml_mul_mat (ctx0, model.per_layer_model_proj, inp_batch);485 per_layer_proj = ggml_scale (ctx0, per_layer_proj, per_layer_projection_scale);486 per_layer_proj = ggml_reshape_3d(ctx0, per_layer_proj, n_embd_per_layer, n_layer, n_tokens);487 488 per_layer_proj = build_norm(per_layer_proj, model.per_layer_proj_norm, nullptr, LLM_NORM_RMS, -1);489 cb(per_layer_proj, "per_layer_proj", -1);490 491 inp_per_layer = ggml_add (ctx0, per_layer_proj, inp_per_layer);492 inp_per_layer = ggml_scale(ctx0, inp_per_layer, per_layer_input_scale);493 cb(inp_per_layer, "inp_per_layer", -1);494 495 // permute to shape: [n_embd_per_layer, n_tokens, n_layer]496 inp_per_layer = ggml_cont(ctx0, ggml_permute(ctx0, inp_per_layer, 0, 2, 1, 3));497 return inp_per_layer;498}499