Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2#include "llama-memory-recurrent.h"3 4void llama_model_qwen35::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_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true);7 8 // Load linear attention (gated delta net) parameters9 ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);10 ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner);11 ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);12 ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);13 ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);14 15 // Mark recurrent layers (linear attention layers). MTP layers are dense16 // attention-only and must be flagged non-recurrent.17 if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {18 uint32_t full_attn_interval = 4;19 ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);20 for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {21 hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);22 }23 }24 25 switch (hparams.n_layer()) {26 case 24: type = hparams.n_embd == 1024 ? LLM_TYPE_0_8B : LLM_TYPE_2B; break;27 case 32: type = hparams.n_embd == 2560 ? LLM_TYPE_4B : LLM_TYPE_9B; break;28 case 64: type = LLM_TYPE_27B; break;29 default: type = LLM_TYPE_UNKNOWN;30 }31}32 33void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {34 LLAMA_LOAD_LOCALS;35 36 const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);37 const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;38 int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;39 40 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);41 42 // output43 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0);44 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);45 46 // if output is NULL, init from the input tok embed47 if (output == NULL) {48 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);49 }50 51 auto load_block_trunk = [&](int il, int flags) {52 auto & layer = layers[il];53 54 // Calculate dimensions from hyperparameters55 const int64_t head_k_dim = hparams.ssm_d_state;56 const int64_t head_v_dim = hparams.ssm_d_state;57 const int64_t n_k_heads = hparams.ssm_n_group;58 const int64_t n_v_heads = hparams.ssm_dt_rank;59 const int64_t key_dim = head_k_dim * n_k_heads;60 const int64_t value_dim = head_v_dim * n_v_heads;61 const int64_t conv_dim = key_dim * 2 + value_dim;62 63 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);64 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags);65 66 if (!hparams.is_recr(il)) {67 // Attention layers68 create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags);69 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags);70 71 // Q/K normalization for attention layers72 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, flags);73 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, flags);74 } else {75 // Linear attention (gated delta net) specific tensors76 // Create tensors with calculated dimensions77 layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);78 layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);79 layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags);80 layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags);81 layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags);82 layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, flags);83 layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, flags);84 layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, flags);85 layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, flags);86 }87 88 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, flags);89 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd}, flags);90 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), {n_embd, n_ff}, flags);91 };92 93 auto load_block_mtp = [&](int il) {94 auto & layer = layers[il];95 96 // MTP block looks like a full-attention Qwen3.5 decoder block.97 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, mtp_flags);98 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, mtp_flags);99 100 create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);101 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, mtp_flags);102 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, mtp_flags);103 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, mtp_flags);104 105 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, mtp_flags);106 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd}, mtp_flags);107 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), {n_embd, n_ff}, mtp_flags);108 109 // NextN-specific tensors that define the MTP block.110 layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);111 layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);112 layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);113 layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);114 layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);115 layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags|TENSOR_NOT_REQUIRED);116 };117 118 for (int i = 0; i < n_layer; ++i) {119 load_block_trunk(i, trunk_flags);120 }121 for (int i = n_layer; i < n_layer_all; ++i) {122 load_block_mtp(i);123 }124}125 126std::unique_ptr<llm_graph_context> llama_model_qwen35::build_arch_graph(const llm_graph_params & params) const {127 if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {128 return std::make_unique<graph_mtp>(*this, params);129 }130 return std::make_unique<graph>(*this, params);131}132 133llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_params & params) :134 llm_build_delta_net_base(params), model(model) {135 const int64_t n_embd_head = hparams.n_embd_head_v();136 137 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());138 139 int sections[4];140 std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);141 142 ggml_tensor * cur;143 ggml_tensor * inpL;144 145 inpL = build_inp_embd(model.tok_embd);146 147 cb(inpL, "model.input_embed", -1);148 149 auto * inp = build_inp_mem_hybrid();150 151 ggml_tensor * inp_pos = build_inp_pos();152 ggml_tensor * inp_out_ids = build_inp_out_ids();153 154 // MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.155 for (int il = 0; il < n_layer; ++il) {156 res->t_layer_inp[il] = inpL;157 158 ggml_tensor * inpSA = inpL;159 160 cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);161 cb(cur, "attn_norm", il);162 163 ggml_build_forward_expand(gf, cur);164 165 // Determine layer type and build appropriate attention mechanism166 if (hparams.is_recr(il)) {167 // Linear attention layer (gated delta net)168 cur = build_layer_attn_linear(inp->get_recr(), cur, il);169 } else {170 // Full attention layer171 cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il);172 }173 174 if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {175 cur = ggml_get_rows(ctx0, cur, inp_out_ids);176 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);177 }178 179 // Residual connection180 cur = ggml_add(ctx0, cur, inpSA);181 cb(cur, "attn_residual", il);182 183 // Save the tensor before post-attention norm for residual connection184 ggml_tensor * ffn_residual = cur;185 186 // Post-attention norm187 ggml_tensor * attn_post_norm = build_norm(cur, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, il);188 cb(attn_post_norm, "attn_post_norm", il);189 190 // Dense FFN layer - without residual connection191 cur = build_layer_ffn(attn_post_norm, il);192 cb(cur, "ffn_out", il);193 194 // Residual connection for FFN - add to the tensor from before post_attention_layernorm195 cur = ggml_add(ctx0, cur, ffn_residual);196 cb(cur, "post_ffn", il);197 198 cur = build_cvec(cur, il);199 cb(cur, "l_out", il);200 201 // Input for next layer202 inpL = cur;203 }204 cur = inpL;205 206 cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);207 208 cb(cur, "h_nextn", -1);209 res->t_h_nextn = cur;210 211 if (!cparams.embeddings_nextn_masked && inp_out_ids) {212 cur = ggml_get_rows(ctx0, cur, inp_out_ids);213 }214 215 cb(cur, "result_norm", -1);216 res->t_embd = cur;217 218 // LM head219 cur = build_lora_mm(model.output, cur, model.output_s);220 221 cb(cur, "result_output", -1);222 res->t_logits = cur;223 224 ggml_build_forward_expand(gf, cur);225}226 227std::pair<ggml_tensor *, ggml_tensor *> llama_model_qwen35::graph::build_qkvz(228 ggml_tensor * input,229 int il) {230 const int64_t n_seqs = ubatch.n_seqs;231 const int64_t n_seq_tokens = ubatch.n_seq_tokens;232 233 ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, input, model.layers[il].wqkv_s);234 qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs);235 cb(qkv_mixed, "linear_attn_qkv_mixed", il);236 237 ggml_tensor * z = build_lora_mm(model.layers[il].wqkv_gate, input, model.layers[il].wqkv_gate_s);238 cb(z, "z", il);239 240 return { qkv_mixed, z };241}242 243ggml_tensor * llama_model_qwen35::graph::build_norm_gated(244 ggml_tensor * input,245 ggml_tensor * weights,246 ggml_tensor * gate,247 int layer) {248 ggml_tensor * normalized = build_norm(input, weights, nullptr, LLM_NORM_RMS, layer);249 ggml_tensor * gated_silu = ggml_silu(ctx0, gate);250 251 return ggml_mul(ctx0, normalized, gated_silu);252}253 254ggml_tensor * llama_model_qwen35::graph::build_layer_attn(255 llm_graph_input_attn_kv * inp,256 ggml_tensor * cur,257 ggml_tensor * inp_pos,258 int * sections,259 int il) {260 const int64_t n_embd_head = hparams.n_embd_head_v();261 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());262 263 // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention264 265 // Qwen3Next uses a single Q projection that outputs query + gate266 auto [Qcur_full, Kcur, Vcur] = build_qkv(model.layers[il], cur,267 n_embd_head * 2, n_head,268 n_embd_head, n_head_kv,269 n_embd_head, n_head_kv,270 il, false);271 cb(Qcur_full, "Qcur_full", il);272 cb(Kcur, "Kcur", il);273 cb(Vcur, "Vcur", il);274 275 ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens,276 ggml_element_size(Qcur_full) * n_embd_head * 2,277 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0);278 cb(Qcur, "Qcur_reshaped", il);279 280 // Apply Q normalization281 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);282 cb(Qcur, "Qcur_normed", il);283 284 // Apply K normalization285 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);286 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);287 cb(Kcur, "Kcur_normed", il);288 289 ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens,290 ggml_element_size(Qcur_full) * n_embd_head * 2,291 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,292 ggml_element_size(Qcur_full) * n_embd_head);293 gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);294 cb(gate, "gate_reshaped", il);295 296 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);297 298 // Apply MRoPE299 Qcur = ggml_rope_multi(300 ctx0, Qcur, inp_pos, nullptr,301 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,302 ext_factor, attn_factor, beta_fast, beta_slow303 );304 305 Kcur = ggml_rope_multi(306 ctx0, Kcur, inp_pos, nullptr,307 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,308 ext_factor, attn_factor, beta_fast, beta_slow309 );310 311 cb(Qcur, "Qcur", il);312 cb(Kcur, "Kcur", il);313 cb(Vcur, "Vcur", il);314 315 // Attention computation316 const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;317 318 cur = build_attn(inp,319 nullptr, nullptr, nullptr,320 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);321 cb(cur, "attn_pregate", il);322 323 ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate);324 cb(gate_sigmoid, "gate_sigmoid", il);325 326 cur = ggml_mul(ctx0, cur, gate_sigmoid);327 cb(cur, "attn_gated", il);328 329 cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);330 cb(cur, "attn_output", il);331 332 return cur;333}334 335ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear(336 llm_graph_input_rs * inp,337 ggml_tensor * cur,338 int il) {339 const auto * mctx_cur = inp->mctx;340 341 const int64_t d_inner = hparams.ssm_d_inner;342 const int64_t n_seqs = ubatch.n_seqs;343 const int64_t head_k_dim = hparams.ssm_d_state;344 const int64_t num_k_heads = hparams.ssm_n_group;345 const int64_t num_v_heads = hparams.ssm_dt_rank;346 const int64_t head_v_dim = d_inner / num_v_heads;347 const int64_t n_seq_tokens = ubatch.n_seq_tokens;348 349 GGML_ASSERT(n_seqs != 0);350 GGML_ASSERT(ubatch.equal_seqs());351 GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);352 353 // Input projections354 auto qkvz = build_qkvz(cur, il);355 ggml_tensor * qkv_mixed = qkvz.first;356 ggml_tensor * z = qkvz.second;357 358 ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s);359 beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs);360 cb(beta, "beta", il);361 362 beta = ggml_sigmoid(ctx0, beta);363 cb(beta, "beta_sigmoid", il);364 365 ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);366 alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);367 cb(alpha, "alpha", il);368 369 ggml_tensor * alpha_biased = ggml_add(ctx0, alpha, model.layers[il].ssm_dt);370 ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased);371 cb(alpha_softplus, "a_softplus", il);372 373 ggml_tensor * gate = ggml_mul(ctx0, alpha_softplus, model.layers[il].ssm_a); // -A_log.exp() * softplus374 cb(gate, "gate", il);375 376 gate = ggml_reshape_4d(ctx0, gate, 1, num_v_heads, n_seq_tokens, n_seqs);377 378 ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);379 ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);380 381 ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d;382 const int64_t conv_kernel_size = conv_kernel->ne[0];383 const int64_t conv_channels = d_inner + 2 * hparams.ssm_n_group * hparams.ssm_d_state;384 385 ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il);386 387 ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs);388 state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs);389 cb(state, "state_predelta", il);390 391 ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel);392 cb(conv_output_proper, "conv_output_raw", il);393 394 ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper);395 cb(conv_output_silu, "conv_output_silu", il);396 397 ggml_tensor * conv_qkv_mix = conv_output_silu;398 399 // Calculate the total conv dimension400 int64_t qkv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads;401 int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, qkv_dim);402 403 // Extract the convolved Q, K, V from conv_output404 ggml_tensor * q_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs,405 ggml_row_size(conv_qkv_mix->type, head_k_dim),406 nb1_qkv,407 nb1_qkv * n_seq_tokens,408 0);409 410 ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs,411 ggml_row_size(conv_qkv_mix->type, head_k_dim),412 nb1_qkv,413 nb1_qkv * n_seq_tokens,414 head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix));415 416 ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs,417 ggml_row_size(conv_qkv_mix->type, head_v_dim),418 nb1_qkv,419 nb1_qkv * n_seq_tokens,420 ggml_row_size(conv_qkv_mix->type, 2 * head_k_dim * num_k_heads));421 422 cb(q_conv, "q_conv", il);423 cb(k_conv, "k_conv", il);424 cb(v_conv, "v_conv", il);425 426 427 const float eps_norm = hparams.f_norm_rms_eps;428 429 q_conv = build_gdn_l2_norm(ctx0, q_conv, eps_norm);430 k_conv = build_gdn_l2_norm(ctx0, k_conv, eps_norm);431 432 //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs);433 //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs);434 //v_conv = ggml_cont_4d(ctx0, v_conv, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);435 436 // if head keys and value keys are different, repeat to force tensors into matching shapes437 // note: need explicit repeat only if we are not using the fused GDN.438 if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) {439 GGML_ASSERT(num_v_heads % num_k_heads == 0);440 q_conv = ggml_repeat_4d(ctx0, q_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs);441 k_conv = ggml_repeat_4d(ctx0, k_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs);442 }443 444 cb(q_conv, "q_conv_predelta", il);445 cb(k_conv, "k_conv_predelta", il);446 cb(v_conv, "v_conv_predelta", il);447 448 ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il);449 450 // z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim]451 ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);452 453 // Apply gated normalization: self.norm(core_attn_out, z)454 ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);455 456 // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim]457 ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs);458 cb(final_output, "final_output", il);459 460 // Output projection461 cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s);462 cb(cur, "linear_attn_out", il);463 464 // Reshape back to original dimensions465 cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs);466 467 return cur;468}469 470ggml_tensor * llama_model_qwen35::graph::build_layer_ffn(ggml_tensor * cur, const int il) {471 // Qwen3.5 does not use MoE FFN472 GGML_ASSERT(model.layers[il].ffn_gate_inp == nullptr);473 474 cur = build_ffn(cur,475 model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s,476 model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,477 model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,478 NULL,479 LLM_FFN_SILU, LLM_FFN_PAR, il);480 cb(cur, "ffn_out", il);481 482 return cur;483}484 485// LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3.5/3.6 dense series486llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)487 : llm_graph_context(params) {488 GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN35 MTP requires n_layer_nextn > 0");489 GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN35 MTP currently only supports a single MTP block");490 491 const int64_t n_embd_head = hparams.n_embd_head_v();492 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());493 494 // hparams.n_layer includes both main model layers and MTP layers. The MTP495 // layer is stored immediately after the main layers in model.layers[].496 const int il = hparams.n_layer();497 const auto & layer = model.layers[il];498 499 GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");500 GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");501 GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");502 503 int sections[4];504 std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);505 506 // TODO: extract in a common llm_graph_context::build_inp_embd_h()507 auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);508 509 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);510 ggml_set_input(inp->tokens);511 512 inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);513 ggml_set_input(inp->embd);514 515 // TODO: make static using `ggml_build_forward_select()`516 // see llm_graph_context::build_inp_embd() for reference517 ggml_tensor * tok_embd;518 if (ubatch.token) {519 ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;520 521 tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);522 } else {523 tok_embd = inp->embd;524 }525 cb(tok_embd, "mtp_tok_embd", il);526 527 inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);528 ggml_set_input(inp->h);529 ggml_set_name(inp->h, "mtp_h_input");530 531 ggml_tensor * h_embd = inp->h;532 533 res->add_input(std::move(inp));534 535 ggml_tensor * inp_pos = build_inp_pos();536 ggml_tensor * inp_out_ids = build_inp_out_ids();537 538 auto * inp_attn = build_attn_inp_kv();539 540 ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);541 cb(h_norm, "mtp_hnorm", il);542 543 ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);544 cb(e_norm, "mtp_enorm", il);545 546 ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);547 cb(concat, "mtp_concat", il);548 549 ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);550 cb(cur, "mtp_eh_proj", il);551 552 ggml_tensor * inpSA = cur;553 554 cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);555 cb(cur, "mtp_attn_norm", il);556 557 auto [Qcur_full, Kcur, Vcur] = build_qkv(layer, cur,558 n_embd_head * 2, n_head,559 n_embd_head, n_head_kv,560 n_embd_head, n_head_kv,561 il, false);562 cb(Qcur_full, "mtp_Qcur_full", il);563 564 ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full,565 n_embd_head, n_head, n_tokens,566 ggml_element_size(Qcur_full) * n_embd_head * 2,567 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,568 0);569 Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);570 cb(Qcur, "mtp_Qcur_normed", il);571 572 ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full,573 n_embd_head, n_head, n_tokens,574 ggml_element_size(Qcur_full) * n_embd_head * 2,575 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,576 ggml_element_size(Qcur_full) * n_embd_head);577 gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);578 cb(gate, "mtp_gate", il);579 580 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);581 Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);582 cb(Kcur, "mtp_Kcur_normed", il);583 584 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);585 cb(Vcur, "mtp_Vcur", il);586 587 Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, nullptr,588 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,589 ext_factor, attn_factor, beta_fast, beta_slow);590 Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr,591 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,592 ext_factor, attn_factor, beta_fast, beta_slow);593 594 const float kq_scale = hparams.f_attention_scale == 0.0f595 ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;596 597 cur = build_attn(inp_attn,598 nullptr, nullptr, nullptr,599 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);600 cb(cur, "mtp_attn_pregate", il);601 602 cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate));603 cur = build_lora_mm(layer.wo, cur, layer.wo_s);604 cb(cur, "mtp_attn_out", il);605 606 cur = ggml_add(ctx0, cur, inpSA);607 cb(cur, "mtp_attn_residual", il);608 609 ggml_tensor * ffn_residual = cur;610 cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);611 cb(cur, "mtp_attn_post_norm", il);612 613 cur = build_ffn(cur,614 layer.ffn_up, nullptr, layer.ffn_up_s,615 layer.ffn_gate, nullptr, layer.ffn_gate_s,616 layer.ffn_down, nullptr, layer.ffn_down_s,617 nullptr,618 LLM_FFN_SILU, LLM_FFN_PAR, il);619 cb(cur, "mtp_ffn_out", il);620 621 cur = ggml_add(ctx0, cur, ffn_residual);622 cb(cur, "mtp_post_ffn", il);623 624 ggml_tensor * head_norm_w = layer.nextn.shared_head_norm625 ? layer.nextn.shared_head_norm626 : model.output_norm;627 GGML_ASSERT(head_norm_w && "QWEN35 MTP: missing both nextn.shared_head_norm and output_norm");628 cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);629 630 cb(cur, "h_nextn", -1);631 res->t_h_nextn = cur;632 633 cur = ggml_get_rows(ctx0, cur, inp_out_ids);634 cb(cur, "mtp_shared_head_norm", -1);635 636 ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;637 ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;638 GGML_ASSERT(head_w && "QWEN35 MTP: missing LM head (nextn.shared_head_head or model.output)");639 cur = build_lora_mm(head_w, cur, head_s);640 cb(cur, "result_output", -1);641 642 res->t_logits = cur;643 ggml_build_forward_expand(gf, cur);644}645