Commit f0c41e016 for llama.cpp
commit f0c41e0168dfd4b5ef72b21d1a311b24cc7a894a
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Tue Oct 6 14:08:54 2026 +0300
models : consolidate nextn row cropping into shared helpers (#30017)
* mimo2 : always emit h_nextn
the other nextn-capable models set it unconditionally
Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD
* models : consolidate nextn row cropping into shared helpers
- replace the duplicated crop conditions and the per-model flags (narrow_early,
crop_before_ffn, crop_last_layer, emit_h_nextn) with two helpers on llm_graph_context:
crop_before_nextn() / crop_after_nextn()
- models that only tested embeddings_nextn_masked now share the same condition, so they
crop the last layer before the nextn capture whenever extraction is off
- t_h_nextn is now set unconditionally in mimo2, qwen4exp and deepseek4 (as in the other
nextn-capable models); host-side reads stay gated by cparams.embeddings_nextn
Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD
diff --git a/src/llama-graph.h b/src/llama-graph.h
index 5cf74c997..838544576 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -1053,6 +1053,16 @@ struct llm_graph_context {
void cb(ggml_tensor * cur, const char * name, int il) const;
+ // true when the last layer must be narrowed to the output rows before the nextn hidden state is captured
+ bool crop_before_nextn(const ggml_tensor * inp_out_ids) const {
+ return inp_out_ids != nullptr && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked);
+ }
+
+ // true when the nextn hidden state must be narrowed to the output rows after it is captured
+ bool crop_after_nextn(const ggml_tensor * inp_out_ids) const {
+ return inp_out_ids != nullptr && cparams.embeddings_nextn && !cparams.embeddings_nextn_masked;
+ }
+
//
// common
//
diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp
index 907b25c67..b04b58449 100644
--- a/src/models/bailingmoe3.cpp
+++ b/src/models/bailingmoe3.cpp
@@ -363,7 +363,7 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
cb(cur, "mla_out", il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -408,7 +408,7 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/cohere2moe.cpp b/src/models/cohere2moe.cpp
index cf2af012d..3baf3950c 100644
--- a/src/models/cohere2moe.cpp
+++ b/src/models/cohere2moe.cpp
@@ -205,7 +205,7 @@ llama_model_cohere2moe::graph::graph(const llama_model & model, const llm_graph_
1.0f / sqrtf(float(n_embd_head)), il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids);
@@ -269,7 +269,7 @@ llama_model_cohere2moe::graph::graph(const llama_model & model, const llm_graph_
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp
index 6d217ea0c..47fc2f1ac 100644
--- a/src/models/deepseek2.cpp
+++ b/src/models/deepseek2.cpp
@@ -634,7 +634,7 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
}
}
- if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -697,7 +697,7 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp
index 849c7a9a0..ade9b9ed9 100644
--- a/src/models/deepseek32.cpp
+++ b/src/models/deepseek32.cpp
@@ -402,7 +402,7 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
}
// when unmasked nextn embeddings are requested, t_h_nextn must keep all rows,
// so the early output masking has to be skipped (it is applied after the final norm instead)
- if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -469,7 +469,7 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp
index 223336d49..4edf0c2b4 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -1335,11 +1335,9 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
ggml_tensor * flat_out = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat;
- if (cparams.embeddings_nextn) {
- ggml_tensor * h_nextn = cparams.embeddings_nextn_masked ? flat_out : inpL;
- cb(h_nextn, "h_nextn", -1);
- res->t_h_nextn = h_nextn;
- }
+ ggml_tensor * h_nextn = cparams.embeddings_nextn_masked ? flat_out : inpL;
+ cb(h_nextn, "h_nextn", -1);
+ res->t_h_nextn = h_nextn;
if (inp_out_ids) {
inpL = ggml_reshape_3d(ctx0, flat_out, n_embd, hc, n_outputs);
diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp
index 38239eba0..fbc4d2a6e 100644
--- a/src/models/gemma4.cpp
+++ b/src/models/gemma4.cpp
@@ -274,7 +274,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
// TODO @ngxson : strip unused token right after the last KV layer to speed up prompt processing
// keep all rows when extracting unmasked nextn embeddings (MTP target needs the hidden state for every token)
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
@@ -374,7 +374,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
ggml_tensor * inp_this_layer = gemma4_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
// TODO @ngxson : improve this
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids);
}
@@ -412,7 +412,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp
index 6a5132cf6..32581f840 100644
--- a/src/models/glm-dsa.cpp
+++ b/src/models/glm-dsa.cpp
@@ -446,7 +446,7 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
}
// when unmasked nextn embeddings are requested, t_h_nextn must keep all rows,
// so the early output masking has to be skipped (it is applied after the final norm instead)
- if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -513,7 +513,7 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/glm4-moe.cpp b/src/models/glm4-moe.cpp
index 8cdbe10ad..d65cace1f 100644
--- a/src/models/glm4-moe.cpp
+++ b/src/models/glm4-moe.cpp
@@ -365,7 +365,7 @@ llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_pa
model.layers[il].wo, NULL, model.layers[il].wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
}
- if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -428,7 +428,7 @@ llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_pa
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index 48f4c1af8..a062deb7a 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -650,8 +650,7 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p
// narrow to the output tokens, then collapse the streams
// Unmasked nextn embeddings need all rows.
- const bool narrow_early = inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked);
- if (narrow_early) {
+ if (crop_before_nextn(inp_out_ids)) {
ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
flat = ggml_get_rows(ctx0, flat, inp_out_ids);
inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs);
@@ -666,7 +665,7 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (inp_out_ids && !narrow_early) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cb(cur, "result_norm", -1);
diff --git a/src/models/hy-v3.cpp b/src/models/hy-v3.cpp
index f6b72d843..bb0e0a2b8 100644
--- a/src/models/hy-v3.cpp
+++ b/src/models/hy-v3.cpp
@@ -148,7 +148,7 @@ llama_model_hy_v3::graph::graph(const llama_model & model, const llm_graph_param
cb(cur, "attn_out", il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -216,7 +216,7 @@ llama_model_hy_v3::graph::graph(const llama_model & model, const llm_graph_param
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp
index ce315f956..866959c83 100644
--- a/src/models/mimo2.cpp
+++ b/src/models/mimo2.cpp
@@ -101,9 +101,7 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
ggml_tensor * inp_out_ids = build_inp_out_ids();
const float v_scale = hparams.f_attn_value_scale;
- const bool emit_h_nextn = cparams.embeddings_nextn;
const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer];
- const bool crop_last_layer = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked) && !extract_final_inp;
for (int il = 0; il < n_layer; ++il) {
res->t_layer_inp[il] = inpL;
@@ -185,7 +183,7 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
}
}
- if (il == n_layer - 1 && crop_last_layer) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids) && !extract_final_inp) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -236,18 +234,17 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
cur = inpL;
if (extract_final_inp) {
res->t_layer_inp[n_layer] = cur;
- if (inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked)) {
+ if (crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
}
- if (emit_h_nextn) {
- cb(cur, "h_nextn", -1);
- res->t_h_nextn = cur;
+ // pre-norm state for MTP; set even when extraction is off
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
- cur = ggml_get_rows(ctx0, cur, inp_out_ids);
- }
+ if (crop_after_nextn(inp_out_ids)) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cur = build_norm(cur,
diff --git a/src/models/nemotron-h-moe.cpp b/src/models/nemotron-h-moe.cpp
index b9b4fdcf1..f1e3ce3b4 100644
--- a/src/models/nemotron-h-moe.cpp
+++ b/src/models/nemotron-h-moe.cpp
@@ -52,11 +52,6 @@ llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, cons
ggml_tensor * inp_out_ids = build_inp_out_ids();
- // attention fills KV over all tokens, but the MoE is position-wise: gather output rows before
- // it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state)
- const bool emit_h_nextn = cparams.embeddings_nextn;
- const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
-
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
@@ -88,7 +83,7 @@ llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, cons
cb(cur, "mtp_attn_residual", il);
// gather the output rows here so the MoE FFN below only runs on the positions we keep
- if (crop_before_ffn) {
+ if (crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
@@ -149,7 +144,7 @@ llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, cons
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!crop_before_ffn && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp
index be27650b0..cd90e24d2 100644
--- a/src/models/nemotron-h.cpp
+++ b/src/models/nemotron-h.cpp
@@ -220,7 +220,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
cur = build_ffn_layer(cur, model, il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked && !extract_final_inp) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids) && !extract_final_inp) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -237,7 +237,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
if (extract_final_inp) {
res->t_layer_inp[n_layer] = cur;
- if (inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
}
@@ -248,7 +248,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/openai-moe.cpp b/src/models/openai-moe.cpp
index 39660a325..a6b63947b 100644
--- a/src/models/openai-moe.cpp
+++ b/src/models/openai-moe.cpp
@@ -114,7 +114,7 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_
cb(cur, "attn_out", il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
// skip computing output for unused tokens
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
@@ -154,7 +154,7 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp
index ab98744a5..e484850da 100644
--- a/src/models/qwen35.cpp
+++ b/src/models/qwen35.cpp
@@ -175,7 +175,7 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -212,7 +212,7 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp
index f0f917af7..6e28028a8 100644
--- a/src/models/qwen35moe.cpp
+++ b/src/models/qwen35moe.cpp
@@ -194,7 +194,7 @@ llama_model_qwen35moe::graph::graph(const llama_model & model, const llm_graph_p
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, sections, il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -232,7 +232,7 @@ llama_model_qwen35moe::graph::graph(const llama_model & model, const llm_graph_p
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp
index 340ef28f7..2d56c38ca 100644
--- a/src/models/qwen3next.cpp
+++ b/src/models/qwen3next.cpp
@@ -168,7 +168,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -206,7 +206,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index 71ce2d503..aca8f6065 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -470,7 +470,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
cur = build_layer_attn(inp->get_attn(), mctx_hyb, inp_kpool, cur, inp_pos, sections, il);
}
- if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
// everything below is per token, so drop the rows that produce no output
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inject = ggml_get_rows(ctx0, inject, inp_out_ids);
@@ -499,13 +499,11 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa
}
// the MTP head reads the hc-wide residual, before the final mixer
- if (cparams.embeddings_nextn) {
- res->t_h_nextn = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]);
- cb(res->t_h_nextn, "h_nextn", -1);
- ggml_build_forward_expand(gf, res->t_h_nextn);
- }
+ res->t_h_nextn = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]);
+ cb(res->t_h_nextn, "h_nextn", -1);
+ ggml_build_forward_expand(gf, res->t_h_nextn);
- if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]);
res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids);
res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]);
diff --git a/src/models/step35.cpp b/src/models/step35.cpp
index ca68855d8..f4660bf19 100644
--- a/src/models/step35.cpp
+++ b/src/models/step35.cpp
@@ -287,7 +287,7 @@ llama_model_step35::graph::graph(const llama_model & model, const llm_graph_para
cb(cur, "attn_proj", il);
}
- if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
+ if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -349,7 +349,7 @@ llama_model_step35::graph::graph(const llama_model & model, const llm_graph_para
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
- if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}