Brunobkr/llama.cpp_AlgMor24_github

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