Enable DeepSeek2 GLM MTP target rows

Author: i386

src/llama-context.cpp

@@ -1938,9 +1938,22 @@ int llama_context::decode(const llama_batch & batch_inp) {
 
                 const uint32_t n_embd  = hparams.n_embd_out();
                 float * embd_nextn_out = embd_nextn.data + offset*n_embd;
+                const size_t n_bytes = (size_t) n_rows*n_embd*sizeof(float);
 
                 GGML_ASSERT((offset + n_rows)*n_embd <= (int64_t) embd_nextn.size);
-                ggml_backend_tensor_get_async(backend_h, t_h_nextn, embd_nextn_out, 0, n_rows*n_embd*sizeof(float));
+                if (n_bytes > ggml_nbytes(t_h_nextn)) {
+                    LLAMA_LOG_ERROR(
+                        "%s: nextn embedding tensor is too small: masked=%d, n_rows=%lld, offset=%lld, "
+                        "n_embd=%u, want=%zu bytes, have=%zu bytes, ne=[%lld, %lld, %lld, %lld], "
+                        "ubatch.n_tokens=%u, n_outputs=%d, n_outputs_prev=%lld, n_tokens_prev=%lld\n",
+                        __func__, (int) masked, (long long) n_rows, (long long) offset,
+                        n_embd, n_bytes, ggml_nbytes(t_h_nextn),
+                        (long long) t_h_nextn->ne[0], (long long) t_h_nextn->ne[1],
+                        (long long) t_h_nextn->ne[2], (long long) t_h_nextn->ne[3],
+                        ubatch.n_tokens, n_outputs, (long long) n_outputs_prev, (long long) n_tokens_prev);
+                    return -3;
+                }
+                ggml_backend_tensor_get_async(backend_h, t_h_nextn, embd_nextn_out, 0, n_bytes);
             }
         }
 

src/models/deepseek2.cpp

@@ -381,7 +381,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) {
+        if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
             cur   = ggml_get_rows(ctx0, cur, inp_out_ids);
             inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
         }
@@ -442,6 +442,10 @@ 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_masked && inp_out_ids) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+
     cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
 
     cb(cur, "result_norm", -1);