Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2#include "llama-memory-recurrent.h"3 4void llama_model_qwen35moe::load_arch_hparams(llama_model_loader & ml) {5 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all, false);6 ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);7 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);8 9 ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true);10 11 // Load linear attention (gated delta net) parameters12 ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);13 ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner);14 ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state);15 ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);16 ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);17 18 // Mark recurrent layers (linear attention layers). MTP layers are dense19 // attention-only and must be flagged non-recurrent.20 if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {21 uint32_t full_attn_interval = 4;22 ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);23 for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {24 hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0);25 }26 }27 28 switch (hparams.n_layer()) {29 case 40: type = LLM_TYPE_35B_A3B; break;30 case 48: type = LLM_TYPE_122B_A10B; break;31 case 60: type = LLM_TYPE_397B_A17B; break;32 default: type = LLM_TYPE_UNKNOWN;33 }34}35 36void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {37 LLAMA_LOAD_LOCALS;38 39 const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);40 const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;41 int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;42 43 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);44 45 // output46 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0);47 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);48 49 // if output is NULL, init from the input tok embed50 if (output == NULL) {51 output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);52 }53 54 auto load_block_trunk = [&](int il, int flags) {55 auto & layer = layers[il];56 57 const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff / n_expert_used;58 const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff;59 60 // Calculate dimensions from hyperparameters61 const int64_t head_k_dim = hparams.ssm_d_state;62 const int64_t head_v_dim = hparams.ssm_d_state;63 const int64_t n_k_heads = hparams.ssm_n_group;64 const int64_t n_v_heads = hparams.ssm_dt_rank;65 const int64_t key_dim = head_k_dim * n_k_heads;66 const int64_t value_dim = head_v_dim * n_v_heads;67 const int64_t conv_dim = key_dim * 2 + value_dim;68 69 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);70 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags);71 72 if (!hparams.is_recr(il)) {73 // Attention layers74 create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags);75 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags);76 77 // Q/K normalization for attention layers78 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, flags);79 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, flags);80 } else {81 // Linear attention (gated delta net) specific tensors82 // Create tensors with calculated dimensions83 layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);84 layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);85 layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags);86 layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags);87 layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags);88 layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, flags);89 layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, flags);90 layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, flags);91 layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, flags);92 }93 94 // Routed experts95 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags);96 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, flags);97 create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, flags);98 99 // Shared experts100 layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, flags);101 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags);102 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags);103 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, flags);104 };105 106 auto load_block_mtp = [&](int il) {107 auto & layer = layers[il];108 109 const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff / n_expert_used;110 const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff;111 112 // MTP block looks like a full-attention Qwen3.5 decoder block with MoE FFN.113 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, mtp_flags);114 layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, mtp_flags);115 116 create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);117 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, mtp_flags);118 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, mtp_flags);119 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, mtp_flags);120 121 // Routed experts122 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, mtp_flags);123 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, mtp_flags);124 create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, mtp_flags);125 126 // Shared experts127 layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, mtp_flags);128 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, mtp_flags);129 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, mtp_flags);130 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, mtp_flags);131 132 // NextN-specific tensors that define the MTP block.133 layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);134 layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);135 layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);136 layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);137 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);138 layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags|TENSOR_NOT_REQUIRED);139 };140 141 for (int i = 0; i < n_layer; ++i) {142 load_block_trunk(i, trunk_flags);143 }144 for (int i = n_layer; i < n_layer_all; ++i) {145 load_block_mtp(i);146 }147}148 149std::unique_ptr<llm_graph_context> llama_model_qwen35moe::build_arch_graph(const llm_graph_params & params) const {150 if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {151 return std::make_unique<graph_mtp>(*this, params);152 }153 return std::make_unique<graph>(*this, params);154}155 156llama_model_qwen35moe::graph::graph(const llama_model & model, const llm_graph_params & params) :157 llm_build_delta_net_base(params), model(model) {158 const int64_t n_embd_head = hparams.n_embd_head_v();159 160 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());161 162 int sections[4];163 std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);164 165 ggml_tensor * cur;166 ggml_tensor * inpL;167 168 inpL = build_inp_embd(model.tok_embd);169 170 cb(inpL, "model.input_embed", -1);171 172 auto * inp = build_inp_mem_hybrid();173 174 ggml_tensor * inp_pos = build_inp_pos();175 ggml_tensor * inp_out_ids = build_inp_out_ids();176 177 // MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.178 for (int il = 0; il < n_layer; ++il) {179 res->t_layer_inp[il] = inpL;180 181 ggml_tensor * inpSA = inpL;182 183 cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);184 cb(cur, "attn_norm", il);185 186 ggml_build_forward_expand(gf, cur);187 188 // Determine layer type and build appropriate attention mechanism189 if (hparams.is_recr(il)) {190 // Linear attention layer (gated delta net)191 cur = build_layer_attn_linear(inp->get_recr(), cur, il);192 } else {193 // Full attention layer194 cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il);195 }196 197 if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {198 cur = ggml_get_rows(ctx0, cur, inp_out_ids);199 inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);200 }201 202 // Residual connection203 cur = ggml_add(ctx0, cur, inpSA);204 cb(cur, "attn_residual", il);205 206 // Save the tensor before post-attention norm for residual connection207 ggml_tensor * ffn_residual = cur;208 209 // Post-attention norm210 ggml_tensor * attn_post_norm = build_norm(cur, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, il);211 cb(attn_post_norm, "attn_post_norm", il);212 213 // MOE FFN layer214 cur = build_layer_ffn(attn_post_norm, il);215 cb(cur, "ffn_out", il);216 217 // Residual connection for FFN - add to the tensor from before post_attention_layernorm218 cur = ggml_add(ctx0, cur, ffn_residual);219 cb(cur, "post_moe", il);220 221 cur = build_cvec(cur, il);222 cb(cur, "l_out", il);223 224 // Input for next layer225 inpL = cur;226 }227 cur = inpL;228 229 // post-norm hidden state feeds both the LM head and the MTP seed below230 cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);231 232 cb(cur, "h_nextn", -1);233 res->t_h_nextn = cur;234 235 if (!cparams.embeddings_nextn_masked && inp_out_ids) {236 cur = ggml_get_rows(ctx0, cur, inp_out_ids);237 }238 239 cb(cur, "result_norm", -1);240 res->t_embd = cur;241 242 // LM head243 cur = build_lora_mm(model.output, cur, model.output_s);244 245 cb(cur, "result_output", -1);246 res->t_logits = cur;247 248 ggml_build_forward_expand(gf, cur);249}250 251std::pair<ggml_tensor *, ggml_tensor *> llama_model_qwen35moe::graph::build_qkvz(252 ggml_tensor * input,253 int il) {254 const int64_t n_seqs = ubatch.n_seqs;255 const int64_t n_seq_tokens = ubatch.n_seq_tokens;256 257 ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, input, model.layers[il].wqkv_s);258 qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs);259 cb(qkv_mixed, "linear_attn_qkv_mixed", il);260 261 ggml_tensor * z = build_lora_mm(model.layers[il].wqkv_gate, input, model.layers[il].wqkv_gate_s);262 cb(z, "z", il);263 264 return { qkv_mixed, z };265}266 267ggml_tensor * llama_model_qwen35moe::graph::build_norm_gated(268 ggml_tensor * input,269 ggml_tensor * weights,270 ggml_tensor * gate,271 int layer) {272 ggml_tensor * normalized = build_norm(input, weights, nullptr, LLM_NORM_RMS, layer);273 ggml_tensor * gated_silu = ggml_silu(ctx0, gate);274 275 return ggml_mul(ctx0, normalized, gated_silu);276}277 278ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn(279 llm_graph_input_attn_kv * inp,280 ggml_tensor * cur,281 ggml_tensor * inp_pos,282 int * sections,283 int il) {284 const int64_t n_embd_head = hparams.n_embd_head_v();285 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());286 287 // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention288 289 // Qwen3Next uses a single Q projection that outputs query + gate290 auto [Qcur_full, Kcur, Vcur] = build_qkv(model.layers[il], cur,291 n_embd_head * 2, n_head,292 n_embd_head, n_head_kv,293 n_embd_head, n_head_kv,294 il, false);295 cb(Qcur_full, "Qcur_full", il);296 cb(Kcur, "Kcur", il);297 cb(Vcur, "Vcur", il);298 299 ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens,300 ggml_element_size(Qcur_full) * n_embd_head * 2,301 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0);302 cb(Qcur, "Qcur_reshaped", il);303 304 // Apply Q normalization305 Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);306 cb(Qcur, "Qcur_normed", il);307 308 // Apply K normalization309 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);310 Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);311 cb(Kcur, "Kcur_normed", il);312 313 ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens,314 ggml_element_size(Qcur_full) * n_embd_head * 2,315 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,316 ggml_element_size(Qcur_full) * n_embd_head);317 gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);318 cb(gate, "gate_reshaped", il);319 320 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);321 322 // Apply IMRoPE323 Qcur = ggml_rope_multi(324 ctx0, Qcur, inp_pos, nullptr,325 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,326 ext_factor, attn_factor, beta_fast, beta_slow327 );328 329 Kcur = ggml_rope_multi(330 ctx0, Kcur, inp_pos, nullptr,331 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,332 ext_factor, attn_factor, beta_fast, beta_slow333 );334 335 cb(Qcur, "Qcur", il);336 cb(Kcur, "Kcur", il);337 cb(Vcur, "Vcur", il);338 339 // Attention computation340 const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;341 342 cur = build_attn(inp,343 nullptr, nullptr, nullptr,344 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);345 cb(cur, "attn_pregate", il);346 347 ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate);348 cb(gate_sigmoid, "gate_sigmoid", il);349 350 cur = ggml_mul(ctx0, cur, gate_sigmoid);351 cb(cur, "attn_gated", il);352 353 cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);354 cb(cur, "attn_output", il);355 356 return cur;357}358 359ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn_linear(360 llm_graph_input_rs * inp,361 ggml_tensor * cur,362 int il) {363 const auto * mctx_cur = inp->mctx;364 365 const int64_t d_inner = hparams.ssm_d_inner;366 const int64_t n_seqs = ubatch.n_seqs;367 const int64_t head_k_dim = hparams.ssm_d_state;368 const int64_t num_k_heads = hparams.ssm_n_group;369 const int64_t num_v_heads = hparams.ssm_dt_rank;370 const int64_t head_v_dim = d_inner / num_v_heads;371 const int64_t n_seq_tokens = ubatch.n_seq_tokens;372 373 GGML_ASSERT(n_seqs != 0);374 GGML_ASSERT(ubatch.equal_seqs());375 GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);376 377 // Input projections378 auto qkvz = build_qkvz(cur, il);379 ggml_tensor * qkv_mixed = qkvz.first;380 ggml_tensor * z = qkvz.second;381 382 ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s);383 beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs);384 cb(beta, "beta", il);385 386 beta = ggml_sigmoid(ctx0, beta);387 cb(beta, "beta_sigmoid", il);388 389 ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s);390 alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs);391 cb(alpha, "alpha", il);392 393 ggml_tensor * alpha_biased = ggml_add(ctx0, alpha, model.layers[il].ssm_dt);394 ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased);395 cb(alpha_softplus, "a_softplus", il);396 397 ggml_tensor * gate = ggml_mul(ctx0, alpha_softplus, model.layers[il].ssm_a); // -A_log.exp() * softplus398 cb(gate, "gate", il);399 400 gate = ggml_reshape_4d(ctx0, gate, 1, num_v_heads, n_seq_tokens, n_seqs);401 402 ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);403 ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il);404 405 ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d;406 const int64_t conv_kernel_size = conv_kernel->ne[0];407 const int64_t conv_channels = d_inner + 2 * hparams.ssm_n_group * hparams.ssm_d_state;408 409 ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il);410 411 ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs);412 state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs);413 cb(state, "state_predelta", il);414 415 ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel);416 cb(conv_output_proper, "conv_output_raw", il);417 418 ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper);419 cb(conv_output_silu, "conv_output_silu", il);420 421 ggml_tensor * conv_qkv_mix = conv_output_silu;422 423 // Calculate the total conv dimension424 int64_t qkv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads;425 int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, qkv_dim);426 427 // Extract the convolved Q, K, V from conv_output428 ggml_tensor * q_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs,429 ggml_row_size(conv_qkv_mix->type, head_k_dim),430 nb1_qkv,431 nb1_qkv * n_seq_tokens,432 0);433 434 ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs,435 ggml_row_size(conv_qkv_mix->type, head_k_dim),436 nb1_qkv,437 nb1_qkv * n_seq_tokens,438 head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix));439 440 ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs,441 ggml_row_size(conv_qkv_mix->type, head_v_dim),442 nb1_qkv,443 nb1_qkv * n_seq_tokens,444 ggml_row_size(conv_qkv_mix->type, 2 * head_k_dim * num_k_heads));445 446 cb(q_conv, "q_conv", il);447 cb(k_conv, "k_conv", il);448 cb(v_conv, "v_conv", il);449 450 451 const float eps_norm = hparams.f_norm_rms_eps;452 453 q_conv = build_gdn_l2_norm(ctx0, q_conv, eps_norm);454 k_conv = build_gdn_l2_norm(ctx0, k_conv, eps_norm);455 456 //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs);457 //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs);458 //v_conv = ggml_cont_4d(ctx0, v_conv, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);459 460 // if head keys and value keys are different, repeat to force tensors into matching shapes461 // note: need explicit repeat only if we are not using the fused GDN.462 if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) {463 GGML_ASSERT(num_v_heads % num_k_heads == 0);464 q_conv = ggml_repeat_4d(ctx0, q_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs);465 k_conv = ggml_repeat_4d(ctx0, k_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs);466 }467 468 cb(q_conv, "q_conv_predelta", il);469 cb(k_conv, "k_conv_predelta", il);470 cb(v_conv, "v_conv_predelta", il);471 472 ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il);473 474 // z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim]475 ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs);476 477 // Apply gated normalization: self.norm(core_attn_out, z)478 ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il);479 480 // Final reshape: [head_dim, n_heads, n_tokens, n_seqs] -> [n_tokens, n_seqs, n_heads * head_dim]481 ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs);482 cb(final_output, "final_output", il);483 484 // Output projection485 cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s);486 cb(cur, "linear_attn_out", il);487 488 // Reshape back to original dimensions489 cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs);490 491 return cur;492}493 494ggml_tensor * llama_model_qwen35moe::graph::build_layer_ffn(ggml_tensor * cur, const int il) {495 // Check if this is an MoE layer496 GGML_ASSERT(model.layers[il].ffn_gate_inp != nullptr);497 498 ggml_tensor * moe_out =499 build_moe_ffn(cur,500 model.layers[il].ffn_gate_inp,501 model.layers[il].ffn_up_exps,502 model.layers[il].ffn_gate_exps,503 model.layers[il].ffn_down_exps,504 nullptr,505 n_expert, n_expert_used,506 LLM_FFN_SILU, true,507 hparams.expert_weights_scale,508 LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,509 nullptr, model.layers[il].ffn_gate_up_exps,510 model.layers[il].ffn_up_exps_s,511 model.layers[il].ffn_gate_exps_s,512 model.layers[il].ffn_down_exps_s);513 cb(moe_out, "ffn_moe_out", il);514 515 // Add shared experts if present - following Qwen3Next reference implementation516 if (model.layers[il].ffn_up_shexp != nullptr) {517 ggml_tensor * ffn_shexp =518 build_ffn(cur,519 model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,520 model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,521 model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,522 NULL,523 LLM_FFN_SILU, LLM_FFN_PAR, il);524 cb(ffn_shexp, "ffn_shexp", il);525 526 // Apply shared expert gating as in the reference implementation527 // The shared expert has its own gate that is sigmoided528 // Note: ffn_gate_inp_shexp is the shared expert gate (outputs 1 value per token)529 ggml_tensor * shared_gate = build_lora_mm(model.layers[il].ffn_gate_inp_shexp, cur);530 cb(shared_gate, "shared_expert_gate", il);531 532 // Apply sigmoid to the gate533 shared_gate = ggml_sigmoid(ctx0, shared_gate);534 cb(shared_gate, "shared_expert_gate_sigmoid", il);535 536 537 // Apply the gate to the shared expert output538 ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate);539 cb(ffn_shexp, "ffn_shexp_gated", il);540 541 cur = ggml_add(ctx0, moe_out, ffn_shexp);542 cb(cur, "ffn_out", il);543 } else {544 cur = moe_out;545 }546 547 return cur;548}549 550// LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3.5/3.6 MoE551llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)552 : llm_graph_context(params) {553 GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN35MOE MTP requires n_layer_nextn > 0");554 GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN35MOE MTP currently only supports a single MTP block");555 556 const int64_t n_embd_head = hparams.n_embd_head_v();557 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());558 559 const int il = hparams.n_layer();560 const auto & layer = model.layers[il];561 562 GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");563 GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");564 GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");565 GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");566 567 int sections[4];568 std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);569 570 // TODO: extract in a common llm_graph_context::build_inp_embd_h()571 auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);572 573 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);574 ggml_set_input(inp->tokens);575 576 inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);577 ggml_set_input(inp->embd);578 579 // TODO: make static using `ggml_build_forward_select()`580 // see llm_graph_context::build_inp_embd() for reference581 ggml_tensor * tok_embd;582 if (ubatch.token) {583 ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;584 585 tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);586 } else {587 tok_embd = inp->embd;588 }589 cb(tok_embd, "mtp_tok_embd", il);590 591 inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);592 ggml_set_input(inp->h);593 ggml_set_name(inp->h, "mtp_h_input");594 595 ggml_tensor * h_embd = inp->h;596 597 res->add_input(std::move(inp));598 599 ggml_tensor * inp_pos = build_inp_pos();600 ggml_tensor * inp_out_ids = build_inp_out_ids();601 602 auto * inp_attn = build_attn_inp_kv();603 604 ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);605 cb(h_norm, "mtp_hnorm", il);606 607 ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);608 cb(e_norm, "mtp_enorm", il);609 610 ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);611 cb(concat, "mtp_concat", il);612 613 ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);614 cb(cur, "mtp_eh_proj", il);615 616 ggml_tensor * inpSA = cur;617 618 cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);619 cb(cur, "mtp_attn_norm", il);620 621 auto [Qcur_full, Kcur, Vcur] = build_qkv(layer, cur,622 n_embd_head * 2, n_head,623 n_embd_head, n_head_kv,624 n_embd_head, n_head_kv,625 il, false);626 cb(Qcur_full, "mtp_Qcur_full", il);627 628 ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full,629 n_embd_head, n_head, n_tokens,630 ggml_element_size(Qcur_full) * n_embd_head * 2,631 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,632 0);633 Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);634 cb(Qcur, "mtp_Qcur_normed", il);635 636 ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full,637 n_embd_head, n_head, n_tokens,638 ggml_element_size(Qcur_full) * n_embd_head * 2,639 ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,640 ggml_element_size(Qcur_full) * n_embd_head);641 gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);642 cb(gate, "mtp_gate", il);643 644 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);645 Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);646 cb(Kcur, "mtp_Kcur_normed", il);647 648 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);649 cb(Vcur, "mtp_Vcur", il);650 651 Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, nullptr,652 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,653 ext_factor, attn_factor, beta_fast, beta_slow);654 Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr,655 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,656 ext_factor, attn_factor, beta_fast, beta_slow);657 658 const float kq_scale = hparams.f_attention_scale == 0.0f659 ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;660 661 cur = build_attn(inp_attn,662 nullptr, nullptr, nullptr,663 Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);664 cb(cur, "mtp_attn_pregate", il);665 666 cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate));667 cur = build_lora_mm(layer.wo, cur, layer.wo_s);668 cb(cur, "mtp_attn_out", il);669 670 cur = ggml_add(ctx0, cur, inpSA);671 cb(cur, "mtp_attn_residual", il);672 673 ggml_tensor * ffn_residual = cur;674 cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);675 cb(cur, "mtp_attn_post_norm", il);676 677 // MoE FFN — routed experts plus gated shared expert (mirrors qwen35moe).678 ggml_tensor * moe_out =679 build_moe_ffn(cur,680 layer.ffn_gate_inp,681 layer.ffn_up_exps,682 layer.ffn_gate_exps,683 layer.ffn_down_exps,684 nullptr,685 n_expert, n_expert_used,686 LLM_FFN_SILU, true,687 hparams.expert_weights_scale,688 LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,689 nullptr, layer.ffn_gate_up_exps,690 layer.ffn_up_exps_s,691 layer.ffn_gate_exps_s,692 layer.ffn_down_exps_s);693 cb(moe_out, "mtp_ffn_moe_out", il);694 695 if (layer.ffn_up_shexp != nullptr) {696 ggml_tensor * ffn_shexp =697 build_ffn(cur,698 layer.ffn_up_shexp, nullptr, layer.ffn_up_shexp_s,699 layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s,700 layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s,701 nullptr,702 LLM_FFN_SILU, LLM_FFN_PAR, il);703 cb(ffn_shexp, "mtp_ffn_shexp", il);704 705 ggml_tensor * shared_gate = build_lora_mm(layer.ffn_gate_inp_shexp, cur);706 shared_gate = ggml_sigmoid(ctx0, shared_gate);707 cb(shared_gate, "mtp_shared_expert_gate_sigmoid", il);708 709 ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate);710 cb(ffn_shexp, "mtp_ffn_shexp_gated", il);711 712 cur = ggml_add(ctx0, moe_out, ffn_shexp);713 } else {714 cur = moe_out;715 }716 cb(cur, "mtp_ffn_out", il);717 718 cur = ggml_add(ctx0, cur, ffn_residual);719 cb(cur, "mtp_post_ffn", il);720 721 ggml_tensor * head_norm_w = layer.nextn.shared_head_norm722 ? layer.nextn.shared_head_norm723 : model.output_norm;724 GGML_ASSERT(head_norm_w && "QWEN35MOE MTP: missing both nextn.shared_head_norm and output_norm");725 cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);726 727 cb(cur, "h_nextn", -1);728 res->t_h_nextn= cur;729 730 cur = ggml_get_rows(ctx0, cur, inp_out_ids);731 cb(cur, "mtp_shared_head_norm", -1);732 733 ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;734 ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;735 GGML_ASSERT(head_w && "QWEN35MOE MTP: missing LM head (nextn.shared_head_head or model.output)");736 cur = build_lora_mm(head_w, cur, head_s);737 cb(cur, "result_output", -1);738 739 res->t_logits = cur;740 ggml_build_forward_expand(gf, cur);741}742