hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Sigbjørn Skjæretmodels : fix incorrect uses of get_key_or_arr (#28868)97e4ca7352026-09-14

1#include "models.h"
2
3void llama_model_step35::load_arch_hparams(llama_model_loader & ml) {
4    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
5
6    hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
7
8    // full_attention layer only use half of the RoPE dimensions
9    hparams.n_rot_full = hparams.n_rot_full / 2;
10
11    // MoE + SWA parameters
12    ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
13    ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
14    ml.get_key(LLM_KV_EXPERT_GATING_FUNC,                hparams.expert_gating_func, false);
15    ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,              hparams.expert_weights_scale, false);
16    ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,               hparams.expert_weights_norm, false);
17
18    // Step35 uses sigmoid gating by default (if not set in GGUF)
19    if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {
20        hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
21    }
22
23    ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW,  hparams.n_swa);
24    ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA,        hparams.rope_freq_base_train_swa, false);
25
26    ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);
27
28    ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP,   hparams.swiglu_clamp_exp,   hparams.n_layer_all, false);
29    ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false);
30
31    switch (hparams.n_layer()) {
32        case 45: type = LLM_TYPE_196B_A11B; break;
33        default: type = LLM_TYPE_UNKNOWN;
34    }
35}
36
37void llama_model_step35::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    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
42    // tensors live in a separate file (e.g. user split target/draft). Mark
43    // MTP tensors NOT_REQUIRED so the trunk loads cleanly.
44    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
45    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
46    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;
47    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;
48
49    if (!ml.load_mtp) {
50        mtp_flags |= TENSOR_SKIP;
51    }
52
53    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
54
55    // output
56    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
57    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, trunk_flags);
58
59    // STEP35 supports per-layer partial RoPE dims; rope factors are stored as a single shared tensor
60    // ("rope_freqs.weight") and ggml uses only the first (n_rot_l/2) entries per layer.
61    uint32_t n_rot_max = 0;
62    for (int i = 0; i < n_layer; ++i) {
63        n_rot_max = std::max(n_rot_max, hparams.n_rot(i));
64    }
65    if (n_rot_max == 0) {
66        n_rot_max = n_rot;
67    }
68
69    auto load_block_trunk = [&](int i, int flags) {
70        auto & layer = layers[i];
71
72        const uint32_t n_head_l      = hparams.n_head(i);
73        const uint32_t n_embd_k_gqa  = hparams.n_embd_k_gqa(i);
74        const uint32_t n_embd_v_gqa  = hparams.n_embd_v_gqa(i);
75
76        layer.attn_norm   = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
77        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
78        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
79
80        // optional rope factors (llama3) / longrope tensors
81        if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) {
82            layer.rope_long  = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG,  "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
83            layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
84        } else {
85            layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));
86        }
87
88        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_l, n_embd_k_gqa, n_embd_v_gqa, flags);
89        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head_l, n_embd}, flags);
90
91        // head-wise attention gate (Step35 self_attn.g_proj)
92        layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head_l}, TENSOR_NOT_REQUIRED);
93
94        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
95
96        // dense MLP (leading dense blocks)
97        layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, TENSOR_NOT_REQUIRED);
98        layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, TENSOR_NOT_REQUIRED);
99        layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, TENSOR_NOT_REQUIRED);
100
101        // MoE routed experts + selection bias (router_bias)
102        const int64_t n_ff_exp = hparams.n_ff_exp();
103        layer.ffn_gate_inp      = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED);
104        layer.ffn_gate_exps     = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp,   n_expert}, TENSOR_NOT_REQUIRED);
105        layer.ffn_down_exps     = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp,   n_embd, n_expert}, TENSOR_NOT_REQUIRED);
106        layer.ffn_up_exps       = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd, n_ff_exp,   n_expert}, TENSOR_NOT_REQUIRED);
107        layer.ffn_exp_probs_b   = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
108
109        // shared expert MLP
110        layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
111        layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
112        layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
113    };
114
115    auto load_block_mtp = [&](int i) {
116        auto & layer = layers[i];
117
118        const uint32_t n_head_l      = hparams.n_head(i);
119        const uint32_t n_embd_k_gqa  = hparams.n_embd_k_gqa(i);
120        const uint32_t n_embd_v_gqa  = hparams.n_embd_v_gqa(i);
121
122        // The MTP block is a full Step3p5 decoder layer (mtp_block) plus the
123        // NextN-specific wiring (enorm/hnorm/eh_proj + optional shared head).
124        // Multi-block MTP: every declared MTP block is required (the draft chain
125        // runs all n_layer_nextn heads), so each block uses the captured
126        // `mtp_flags` directly — already NOT_REQUIRED for a trunk-only GGUF,
127        // which keeps that path correct.
128
129        layer.attn_norm   = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, mtp_flags);
130        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
131        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, TENSOR_NOT_REQUIRED);
132
133        if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) {
134            layer.rope_long  = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG,  "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | TENSOR_DUPLICATED);
135            layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | TENSOR_DUPLICATED);
136        } else {
137            layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot_max/2}, TENSOR_NOT_REQUIRED | TENSOR_DUPLICATED);
138        }
139
140        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_l, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);
141        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_v * n_head_l, n_embd}, mtp_flags);
142
143        layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head_l}, TENSOR_NOT_REQUIRED);
144
145        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, mtp_flags);
146
147        // dense MLP (leading dense blocks) — present if the MTP block isn't MoE
148        layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, TENSOR_NOT_REQUIRED);
149        layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, TENSOR_NOT_REQUIRED);
150        layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, TENSOR_NOT_REQUIRED);
151
152        // MoE routed experts + selection bias (router_bias)
153        const int64_t n_ff_exp = hparams.n_ff_exp();
154        layer.ffn_gate_inp      = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), {n_embd, n_expert}, TENSOR_NOT_REQUIRED);
155        layer.ffn_gate_exps     = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp,   n_expert}, TENSOR_NOT_REQUIRED);
156        layer.ffn_down_exps     = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp,   n_embd, n_expert}, TENSOR_NOT_REQUIRED);
157        layer.ffn_up_exps       = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd, n_ff_exp,   n_expert}, TENSOR_NOT_REQUIRED);
158        layer.ffn_exp_probs_b   = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);
159
160        layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
161        layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, hparams.n_ff_shexp}, TENSOR_NOT_REQUIRED);
162        layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED);
163
164        // NextN-specific tensors that define the MTP block.
165        layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", i), { 2 * n_embd, n_embd }, mtp_flags);
166        layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", i), { n_embd },              mtp_flags);
167        layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", i), { n_embd },              mtp_flags);
168        layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS,     "weight", i), { n_embd, n_vocab },     TENSOR_NOT_REQUIRED);
169        layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab },     TENSOR_NOT_REQUIRED);
170        layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd },              TENSOR_NOT_REQUIRED);
171    };
172
173    for (int i = 0; i < n_layer; ++i) {
174        load_block_trunk(i, trunk_flags);
175    }
176    // All n_layer_nextn MTP blocks are required — the multi-block draft chain
177    // runs every head (head k at offset k). The GGUF declares the count via
178    // step35.nextn_predict_layers.
179    for (int i = n_layer; i < n_layer_all; ++i) {
180        load_block_mtp(i);
181    }
182}
183
184std::unique_ptr<llm_graph_context> llama_model_step35::build_arch_graph(const llm_graph_params & params) const {
185    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
186        return std::make_unique<graph_mtp>(*this, params);
187    }
188    return std::make_unique<graph>(*this, params);
189}
190
191llama_model_step35::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
192    ggml_tensor * cur;
193    ggml_tensor * inpL;
194
195    inpL = build_inp_embd(model.tok_embd);
196    ggml_tensor * inp_pos     = build_inp_pos();
197    auto        * inp_attn    = build_attn_inp_kv_iswa();
198    ggml_tensor * inp_out_ids = build_inp_out_ids();
199
200    // MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.
201    for (int il = 0; il < n_layer; ++il) {
202        ggml_tensor * inpSA = inpL;
203
204        const uint32_t n_head_l    = hparams.n_head(il);
205        const uint32_t n_head_kv_l = hparams.n_head_kv(il);
206
207        const float freq_base_l  = model.get_rope_freq_base(cparams, il);
208        const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
209
210        cur = inpL;
211
212        // dump pre-attn RMSNorm input to pinpoint layer boundary issues
213        cb(cur, "attn_norm_in", il);
214
215        // self-attention
216        {
217            cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);
218            cb(cur, "attn_norm", il);
219            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
220                    n_embd_head_k, n_head_l,
221                    n_embd_head_k, n_head_kv_l,
222                    n_embd_head_v, n_head_kv_l,
223                    il, false);
224
225            cb(Qcur, "Qcur", il);
226            cb(Kcur, "Kcur", il);
227            cb(Vcur, "Vcur", il);
228
229            Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l,    n_tokens);
230            Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens);
231            Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens);
232
233            // Q/K per-head RMSNorm (Step35 q_norm / k_norm)
234            if (model.layers[il].attn_q_norm) {
235                Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);
236                cb(Qcur, "Qcur_normed", il);
237            }
238            if (model.layers[il].attn_k_norm) {
239                Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il);
240                cb(Kcur, "Kcur_normed", il);
241            }
242
243            // RoPE (partial rotary factors per layer)
244            const bool is_swa = hparams.is_swa(il);
245            ggml_tensor * rope_factors = is_swa ? nullptr : model.get_rope_factors(cparams, il);
246            const int64_t n_rot_l = hparams.n_rot(il);
247            Qcur = ggml_rope_ext(
248                ctx0, Qcur, inp_pos, rope_factors,
249                n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
250                ext_factor, attn_factor, beta_fast, beta_slow
251            );
252            Kcur = ggml_rope_ext(
253                ctx0, Kcur, inp_pos, rope_factors,
254                n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
255                ext_factor, attn_factor, beta_fast, beta_slow
256            );
257            cb(Qcur, "Qcur_pos", il);
258            cb(Kcur, "Kcur_pos", il);
259
260            const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));
261            ggml_tensor * attn_out = build_attn(inp_attn,
262                    nullptr, nullptr, nullptr,
263                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
264            cb(attn_out, "attn_out", il);
265            // head-wise attention gate: sigmoid(g_proj(x)) in torch
266            if (model.layers[il].wqkv_gate) {
267                ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, cur); // [n_head_l, n_tokens]
268                cb(gate, "attn_gate", il);
269
270                gate = ggml_sigmoid(ctx0, gate);
271                cb(gate, "attn_gate_sigmoid", il);
272
273                // reshape + broadcast to [n_embd_head_v, n_head_l, n_tokens]
274                ggml_tensor * attn_3d = ggml_reshape_3d(ctx0, attn_out, n_embd_head_v, n_head_l, n_tokens);
275                ggml_tensor * gate_3d = ggml_reshape_3d(ctx0, gate,       1,          n_head_l, n_tokens);
276                cb(gate_3d, "attn_gate_3d", il);
277
278                attn_3d = ggml_mul(ctx0, attn_3d, gate_3d);
279                cb(attn_3d, "attn_gated_3d", il);
280
281                attn_out = ggml_reshape_2d(ctx0, attn_3d, n_embd_head_v * n_head_l, n_tokens);
282                cb(attn_out, "attn_gated", il);
283            }
284
285            // output projection
286            cur = build_lora_mm(model.layers[il].wo, attn_out, model.layers[il].wo_s);
287            cb(cur, "attn_proj", il);
288        }
289
290        if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
291            cur   = ggml_get_rows(ctx0, cur, inp_out_ids);
292            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
293        }
294
295        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
296        cb(ffn_inp, "ffn_inp", il);
297
298        cur = build_norm(ffn_inp, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il);
299        cb(cur, "ffn_norm", il);
300
301        // feed-forward
302        if (model.layers[il].ffn_gate_inp == nullptr) {
303            // dense MLP
304            cur = build_ffn(cur,
305                    model.layers[il].ffn_up,   model.layers[il].ffn_up_b,   nullptr,
306                    model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, nullptr,
307                    model.layers[il].ffn_down, model.layers[il].ffn_down_b, nullptr,
308                    nullptr,
309                    LLM_FFN_SILU, LLM_FFN_PAR, il);
310            cb(cur, "ffn_out", il);
311        } else {
312            // MoE routed experts
313            ggml_tensor * moe_out = build_moe_ffn(cur,
314                    model.layers[il].ffn_gate_inp,
315                    model.layers[il].ffn_up_exps,
316                    model.layers[il].ffn_gate_exps,
317                    model.layers[il].ffn_down_exps,
318                    model.layers[il].ffn_exp_probs_b,
319                    n_expert, n_expert_used,
320                    LLM_FFN_SILU, hparams.expert_weights_norm,
321                    hparams.expert_weights_scale,
322                    (llama_expert_gating_func_type) hparams.expert_gating_func,
323                    il);
324            cb(moe_out, "ffn_moe_out", il);
325
326            // shared expert MLP (always added on MoE layers in Step35)
327            ggml_tensor * sh_out = build_ffn(cur,
328                    model.layers[il].ffn_up_shexp,   nullptr, nullptr,
329                    model.layers[il].ffn_gate_shexp, nullptr, nullptr,
330                    model.layers[il].ffn_down_shexp, nullptr, nullptr,
331                    nullptr,
332                    LLM_FFN_SILU, LLM_FFN_PAR, il);
333            cb(sh_out, "ffn_shared_out", il);
334
335            cur = ggml_add(ctx0, moe_out, sh_out);
336            cb(cur, "ffn_out", il);
337        }
338        cur = ggml_add(ctx0, cur, ffn_inp);
339
340        cur = build_cvec(cur, il);
341        cb(cur, "l_out", il);
342
343        // input for next layer
344        inpL = cur;
345    }
346
347    cur = inpL;
348
349    cb(cur, "h_nextn", -1);
350    res->t_h_nextn = cur;
351
352    if (!cparams.embeddings_nextn_masked && inp_out_ids) {
353        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
354    }
355
356    cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
357    cb(cur, "result_norm", -1);
358    res->t_embd = cur;
359
360    cur = build_lora_mm(model.output, cur, model.output_s);
361    cb(cur, "result_output", -1);
362    res->t_logits = cur;
363
364    ggml_build_forward_expand(gf, cur);
365}
366
367// LLM_GRAPH_TYPE_DECODER_MTP draft head for Step3p5 (MoE)
368llama_model_step35::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
369    : llm_graph_context(params) {
370    GGML_ASSERT(hparams.n_layer_nextn > 0 && "STEP35 MTP requires n_layer_nextn > 0");
371
372    // Multi-block MTP: the DECODER_MTP graph runs the MTP head selected by
373    // cparams.nextn_layer_offset (0 = first trained head). The speculative driver
374    // bumps the offset per draft step to chain heads 45->46->47. offset 0 keeps
375    // single-block behavior identical to before.
376    const int il = hparams.n_layer() + cparams.nextn_layer_offset;
377    GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
378                cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
379                "nextn_layer_offset out of range [0, n_layer_nextn)");
380    const auto & layer = model.layers[il];
381
382    GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
383    GGML_ASSERT(layer.nextn.enorm   && "MTP block missing nextn.enorm");
384    GGML_ASSERT(layer.nextn.hnorm   && "MTP block missing nextn.hnorm");
385
386    const uint32_t n_head_l    = hparams.n_head(il);
387    const uint32_t n_head_kv_l = hparams.n_head_kv(il);
388
389    const float freq_base_l  = model.get_rope_freq_base(cparams, il);
390    const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
391
392    auto inp = std::make_unique<llm_graph_input_embd>(hparams.n_embd);
393
394    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
395    ggml_set_input(inp->tokens);
396
397    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
398    ggml_set_input(inp->embd);
399    ggml_set_name(inp->embd, "mtp_h_input");
400
401    ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
402
403    ggml_tensor * h_input  = inp->embd;
404    ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
405    cb(tok_embd, "mtp_tok_embd", il);
406
407    res->add_input(std::move(inp));
408
409    ggml_tensor * inp_pos  = build_inp_pos();
410    auto        * inp_attn = build_attn_inp_kv_iswa();
411
412    ggml_tensor * h_norm = build_norm(h_input, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
413    cb(h_norm, "mtp_hnorm", il);
414
415    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
416    cb(e_norm, "mtp_enorm", il);
417
418    ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
419    cb(concat, "mtp_concat", il);
420
421    ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat);
422    cb(cur, "mtp_eh_proj", il);
423
424    ggml_tensor * inpSA = cur;
425
426    // mtp_block: full Step3p5 decoder layer (attention with optional head-wise gate, then MoE/dense FFN)
427    cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
428    cb(cur, "mtp_attn_norm", il);
429
430    auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur,
431            n_embd_head_k, n_head_l,
432            n_embd_head_k, n_head_kv_l,
433            n_embd_head_v, n_head_kv_l,
434            il, false);
435    cb(Qcur, "mtp_Qcur", il);
436    cb(Kcur, "mtp_Kcur", il);
437    cb(Vcur, "mtp_Vcur", il);
438
439    Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k, n_head_l,    n_tokens);
440    Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_kv_l, n_tokens);
441    Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_kv_l, n_tokens);
442
443    if (layer.attn_q_norm) {
444        Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);
445        cb(Qcur, "mtp_Qcur_normed", il);
446    }
447    if (layer.attn_k_norm) {
448        Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);
449        cb(Kcur, "mtp_Kcur_normed", il);
450    }
451
452    const bool    is_swa       = hparams.is_swa(il);
453    ggml_tensor * rope_factors = is_swa ? nullptr : model.get_rope_factors(cparams, il);
454    const int64_t n_rot_l      = hparams.n_rot(il);
455
456    Qcur = ggml_rope_ext(
457        ctx0, Qcur, inp_pos, rope_factors,
458        n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
459        ext_factor, attn_factor, beta_fast, beta_slow);
460    Kcur = ggml_rope_ext(
461        ctx0, Kcur, inp_pos, rope_factors,
462        n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
463        ext_factor, attn_factor, beta_fast, beta_slow);
464    cb(Qcur, "mtp_Qcur_pos", il);
465    cb(Kcur, "mtp_Kcur_pos", il);
466
467    const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));
468    ggml_tensor * attn_out = build_attn(inp_attn,
469            nullptr, nullptr, nullptr,
470            Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
471    cb(attn_out, "mtp_attn_out", il);
472
473    // head-wise attention gate: sigmoid(g_proj(x))
474    if (layer.wqkv_gate) {
475        ggml_tensor * gate = build_lora_mm(layer.wqkv_gate, cur); // [n_head_l, n_tokens]
476        cb(gate, "mtp_attn_gate", il);
477
478        gate = ggml_sigmoid(ctx0, gate);
479        cb(gate, "mtp_attn_gate_sigmoid", il);
480
481        ggml_tensor * attn_3d = ggml_reshape_3d(ctx0, attn_out, n_embd_head_v, n_head_l, n_tokens);
482        ggml_tensor * gate_3d = ggml_reshape_3d(ctx0, gate,       1,           n_head_l, n_tokens);
483        cb(gate_3d, "mtp_attn_gate_3d", il);
484
485        attn_3d = ggml_mul(ctx0, attn_3d, gate_3d);
486        cb(attn_3d, "mtp_attn_gated_3d", il);
487
488        attn_out = ggml_reshape_2d(ctx0, attn_3d, n_embd_head_v * n_head_l, n_tokens);
489        cb(attn_out, "mtp_attn_gated", il);
490    }
491
492    cur = build_lora_mm(layer.wo, attn_out, layer.wo_s);
493    cb(cur, "mtp_attn_proj", il);
494
495    cur = ggml_add(ctx0, cur, inpSA);
496    cb(cur, "mtp_attn_residual", il);
497
498    ggml_tensor * ffn_inp = cur;
499    cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
500    cb(cur, "mtp_ffn_norm", il);
501
502    // FFN: dense MLP or MoE (mirrors trunk path)
503    if (layer.ffn_gate_inp == nullptr) {
504        cur = build_ffn(cur,
505                layer.ffn_up,   layer.ffn_up_b,   nullptr,
506                layer.ffn_gate, layer.ffn_gate_b, nullptr,
507                layer.ffn_down, layer.ffn_down_b, nullptr,
508                nullptr,
509                LLM_FFN_SILU, LLM_FFN_PAR, il);
510        cb(cur, "mtp_ffn_out", il);
511    } else {
512        ggml_tensor * moe_out = build_moe_ffn(cur,
513                layer.ffn_gate_inp,
514                layer.ffn_up_exps,
515                layer.ffn_gate_exps,
516                layer.ffn_down_exps,
517                layer.ffn_exp_probs_b,
518                n_expert, n_expert_used,
519                LLM_FFN_SILU, hparams.expert_weights_norm,
520                hparams.expert_weights_scale,
521                (llama_expert_gating_func_type) hparams.expert_gating_func,
522                il);
523        cb(moe_out, "mtp_ffn_moe_out", il);
524
525        ggml_tensor * sh_out = build_ffn(cur,
526                layer.ffn_up_shexp,   nullptr, nullptr,
527                layer.ffn_gate_shexp, nullptr, nullptr,
528                layer.ffn_down_shexp, nullptr, nullptr,
529                nullptr,
530                LLM_FFN_SILU, LLM_FFN_PAR, il);
531        cb(sh_out, "mtp_ffn_shared_out", il);
532
533        cur = ggml_add(ctx0, moe_out, sh_out);
534        cb(cur, "mtp_ffn_out", il);
535    }
536    cur = ggml_add(ctx0, cur, ffn_inp);
537    cb(cur, "mtp_post_ffn", il);
538
539    ggml_tensor * inp_out_ids = build_inp_out_ids();
540    cur = ggml_get_rows(ctx0, cur, inp_out_ids);
541
542    // Pre-norm hidden state: used by the AR draft loop to seed the next MTP step.
543    cb(cur, "h_nextn", -1);
544    res->t_h_nextn = cur;
545
546    ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
547            ? layer.nextn.shared_head_norm
548            : model.output_norm;
549    GGML_ASSERT(head_norm_w && "STEP35 MTP: missing both nextn.shared_head_norm and output_norm");
550    cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
551    cb(cur, "mtp_shared_head_norm", -1);
552
553    ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
554    GGML_ASSERT(head_w && "STEP35 MTP: missing LM head (nextn.shared_head_head or model.output)");
555    cur = build_lora_mm(head_w, cur);
556    cb(cur, "result_output", -1);
557
558    res->t_logits = cur;
559    ggml_build_forward_expand(gf, cur);
560}