Commit a8fad21e9 for llama.cpp
commit a8fad21e9949c20f6b21692c8d38a8aba422eafe
Author: Aman Gupta <amangupta052@gmail.com>
Date: Mon Nov 9 20:52:32 2026 +0800
mtp: read the target hidden states from a device buffer (#31927)
diff --git a/common/common.cpp b/common/common.cpp
index 75fd39665..cca920f09 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -2242,6 +2242,14 @@ bool common_batch::set_embd_state(int32_t idx, llama_embd state) {
return true;
}
+bool common_batch::set_embd_state_row(int32_t idx, int32_t row) {
+ if (idx < 0 || idx >= size() || row < 0 || tokens[idx].state.data != nullptr || tokens[idx].state_row >= 0) {
+ return false;
+ }
+ tokens[idx].state_row = row;
+ return true;
+}
+
int32_t common_batch::add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output) {
token t = { LLAMA_TOKEN_NULL, { 0, 0, 0, 0 }, seq_id, output, embd, { nullptr, 0, 0 }, {} };
for (int32_t j = 0; j < n_pos; ++j) {
@@ -2307,6 +2315,9 @@ llama_batch_ext * common_batch::get_sub_batch(int32_t off, int32_t n) {
if (t.state.data) {
llama_batch_ext_set_embd_state(res, idx, t.state); // contexts without a state input ignore it
}
+ if (t.state_row >= 0 && !llama_batch_ext_set_embd_state_row(res, idx, t.state_row)) {
+ GGML_ABORT("%s: failed to set the state row of the entry at index %d\n", __func__, i);
+ }
if (t.decision_order != 0) {
llama_batch_ext_set_decision_order(res, idx, (llama_decision_order) t.decision_order);
}
diff --git a/common/common.h b/common/common.h
index 84abddc76..3aae7a8da 100644
--- a/common/common.h
+++ b/common/common.h
@@ -1094,6 +1094,7 @@ struct common_batch {
std::vector<llama_seq_id> seq_ids_extra; // see add_seq()
int32_t decision_order = 0; // see llama_batch_ext_set_decision_order()
int32_t embd_row = -1; // see add_embd_row(), -1 if none
+ int32_t state_row = -1; // see set_embd_state_row(), -1 if none
};
std::vector<token> tokens; // mirror of the entries, tokens[i] describes batch index i
@@ -1132,6 +1133,9 @@ struct common_batch {
// attach a state embedding (e.g. the target hidden state for MTP) to the entry at idx, can only be set once per entry
bool set_embd_state(int32_t idx, llama_embd state);
+ // the state of entry idx is row `row` of the h_nextn device buffer of the state source context, see llama_set_embd_state_src()
+ bool set_embd_state_row(int32_t idx, int32_t row);
+
// add an embedding-only entry (no token id)
// pos points to n_pos positions
int32_t add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output);
diff --git a/common/speculative.cpp b/common/speculative.cpp
index e6290eaad..6b5c9bc36 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -1563,6 +1563,13 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
std::vector<int> i_last;
std::vector<std::vector<float>> chain_h;
+ // target h rows on the device (llama_set_embeddings_nextn_dev): the catch-up in process() reads row k - 1 of the last
+ // target decode in place, and row n_batch_tgt + seq_id for the first token of seq_id, uploaded from pending_h
+ // pending_h stays on the host: process() and accept() fetch it before the next target decode overwrites the rows
+ bool dev_rows = false;
+ int32_t n_batch_tgt = 0;
+ std::vector<int32_t> verify_row0; // [n_seq] device row of the first verify row, see verify_h_rows
+
common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq)
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq, params.draft.n_max)
, params(params.draft)
@@ -1637,10 +1644,27 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
verify_h.assign(n_seq, {});
verify_h_rows.assign(n_seq, 0);
+
+ n_batch_tgt = (int32_t) llama_n_batch(ctx_tgt);
+ verify_row0.assign(n_seq, -1);
+
+ // shared memory and chained heads keep the host rows
+ if (!is_mem_shared && !chain_heads) {
+ dev_rows = llama_set_embeddings_nextn_dev(ctx_tgt, true) && llama_set_embd_state_src(ctx_dft, ctx_tgt);
+ if (!dev_rows) {
+ llama_set_embd_state_src(ctx_dft, nullptr);
+ llama_set_embeddings_nextn_dev(ctx_tgt, false);
+ }
+ }
+ SPC_TRC("- target h rows on the device: %s\n", dev_rows ? "yes" : "no");
}
~common_speculative_impl_draft_mtp() override {
auto * ctx_dft = this->params.ctx_dft;
+ if (dev_rows) {
+ llama_set_embd_state_src(ctx_dft, nullptr);
+ llama_set_embeddings_nextn_dev(this->params.ctx_tgt, false);
+ }
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) {
if (backend_chains[seq_id] == nullptr) {
continue;
@@ -1715,7 +1739,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
// i.e. we cannot have seq_id like this: [0, 0, 0, 1, 1, 0, 1, 1]
// ^--- this is a problem
// TODO:this is generally true, but would be nice to assert it
- const float * h_tgt = llama_get_embeddings_nextn(ctx_tgt);
+ if (dev_rows) {
+ GGML_ASSERT(n_tokens <= n_batch_tgt && "the target batch exceeds the h_nextn device buffer");
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ if (i_batch_beg[seq_id] >= 0) {
+ llama_set_embeddings_nextn_rows(ctx_tgt, n_batch_tgt + seq_id, 1, pending_h[seq_id].data());
+ }
+ }
+ }
+
+ const float * h_tgt = dev_rows ? nullptr : llama_get_embeddings_nextn(ctx_tgt);
for (int k = 0; k < n_tokens; ++k) {
const auto & t = batch_in.tokens[k];
@@ -1727,6 +1760,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
? batch.add(t.id, t.pos[0], seq_id, false)
: batch.add_embd(t.embd, t.pos.data(), seq_id, false);
+ if (dev_rows) {
+ batch.set_embd_state_row(idx, k == i_batch_beg[seq_id] ? n_batch_tgt + seq_id : k - 1);
+ continue;
+ }
+
const float * h_row = k == i_batch_beg[seq_id]
? pending_h[seq_id].data()
: h_tgt + (size_t) (k - 1) * n_embd;
@@ -1785,6 +1823,14 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
const int32_t h_offset = is_mem_shared && !all_rows_output ? i_batch_end[seq_id] : i_batch_beg[seq_id];
verify_h_rows[seq_id] = h_rows;
+
+ if (dev_rows) {
+ // the rows stay on the device until the next target decode, only the pending row comes to the host
+ verify_row0[seq_id] = h_offset;
+ llama_get_embeddings_nextn_rows(ctx_tgt, h_offset + h_rows - 1, 1, pending_h[seq_id].data());
+ continue;
+ }
+
verify_h[seq_id].resize((size_t) h_rows * n_embd);
for (int32_t i = 0; i < h_rows; ++i) {
@@ -1976,6 +2022,13 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
}
const int32_t i_h = std::min<int32_t>(n_accepted, n_rows - 1);
+ if (dev_rows) {
+ // process() already fetched the last row
+ if (i_h != n_rows - 1) {
+ llama_get_embeddings_nextn_rows(params.ctx_tgt, verify_row0[seq_id] + i_h, 1, pending_h[seq_id].data());
+ }
+ return;
+ }
const size_t row_bytes = (size_t) n_embd * sizeof(float);
std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes);
}
diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index f84c6d50b..7fce6624f 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -59,8 +59,14 @@ bool llama_batch_allocr::init(
int32_t n_embd_only = 0;
int32_t n_both = 0;
- const bool has_state = batch_inp.tokens[0].has_state;
- const bool has_row = batch_inp.tokens[0].embd_row >= 0;
+ const bool has_state = batch_inp.tokens[0].has_state;
+ const bool has_row = batch_inp.tokens[0].embd_row >= 0;
+ const bool has_state_row = batch_inp.tokens[0].state_row >= 0;
+
+ if (has_state && has_state_row) {
+ LLAMA_LOG_ERROR("%s: the entries have a state embedding and a state row\n", __func__);
+ return false;
+ }
for (int32_t i = 0; i < n_tok; ++i) {
const bool is_tok = batch_inp.tokens[i].id != LLAMA_TOKEN_NULL;
@@ -87,6 +93,11 @@ bool llama_batch_allocr::init(
return false;
}
+ if ((batch_inp.tokens[i].state_row >= 0) != has_state_row) {
+ LLAMA_LOG_ERROR("%s: all entries in the batch must have the same state row presence\n", __func__);
+ return false;
+ }
+
n_tok_only += is_tok && !is_emb;
n_embd_only += is_emb && !is_tok;
n_both += is_tok && is_emb;
@@ -157,6 +168,13 @@ bool llama_batch_allocr::init(
}
}
+ if (has_state_row) {
+ state_row_vec.resize(n_tok);
+ for (int32_t i = 0; i < n_tok; ++i) {
+ state_row_vec[i] = batch_inp.tokens[i].state_row;
+ }
+ }
+
//
// build flat pos array, section-major: pos[j*n_tok + i] = section j of entry i
// token entry: [p, p, p, 0] (M-RoPE text position)
@@ -861,6 +879,7 @@ void llama_batch_allocr::clear() {
is_embd_vec .clear();
state_vec .clear();
embd_row_vec.clear();
+ state_row_vec.clear();
seq_id_data .clear();
pos .clear();
n_seq_id .clear();
@@ -905,6 +924,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
const bool use_embd = batch.embd && !(mixed_batch && n_embd_rows == 0);
const bool has_state = !state_vec.empty();
const bool has_row = !embd_row_vec.empty();
+ const bool has_srow = !state_row_vec.empty();
const int64_t n_embd_all = use_embd ? (int64_t) n_tokens*n_embd : 0;
const int64_t n_state_all = has_state ? (int64_t) n_tokens*n_embd_state : 0;
@@ -914,6 +934,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
udata->embd .resize(n_embd_all);
udata->embd_state.resize(n_state_all);
udata->embd_row .resize(has_row ? n_tokens : 0);
+ udata->embd_state_row.resize(has_srow ? n_tokens : 0);
udata->pos .resize(n_pos_all);
udata->n_seq_id .resize(n_tokens);
udata->seq_id .resize(n_tokens);
@@ -949,6 +970,10 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
udata->embd_row[i] = embd_row_vec[idxs[i]];
}
+ if (has_srow) {
+ udata->embd_state_row[i] = state_row_vec[idxs[i]];
+ }
+
for (size_t j = 0; j < (size_t)n_pos_per_embd; ++j) {
udata->pos[j*n_tokens + i] = batch.pos[j*batch.n_tokens + idxs[i]];
}
@@ -1354,6 +1379,23 @@ bool llama_batch_ext::set_token_embd_row(int32_t idx, int32_t row) {
return true;
}
+bool llama_batch_ext::set_token_state_row(int32_t idx, int32_t row) {
+ if (idx < 0 || idx >= (int32_t) tokens.size() || row < 0) {
+ return false;
+ }
+
+ token & t = tokens[idx];
+
+ if (t.has_state || t.state_row >= 0) {
+ LLAMA_LOG_ERROR("%s: entry %d already has a state embedding or a state row\n", __func__, idx);
+ return false;
+ }
+
+ t.state_row = row;
+
+ return true;
+}
+
bool llama_batch_ext::set_token_pos(int32_t idx, const llama_pos * pos_in) {
if (idx < 0 || idx >= (int32_t) tokens.size()) {
return false;
@@ -1455,6 +1497,10 @@ bool llama_batch_ext_set_embd_state(llama_batch_ext * batch, int32_t idx, llama_
return batch->set_token_state(idx, embd);
}
+bool llama_batch_ext_set_embd_state_row(llama_batch_ext * batch, int32_t idx, int32_t row) {
+ return batch->set_token_state_row(idx, row);
+}
+
bool llama_batch_ext_set_output_embd(llama_batch_ext * batch, int32_t idx, bool value) {
return batch->set_output(idx, value);
}
diff --git a/src/llama-batch.h b/src/llama-batch.h
index d2361396c..d16662898 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -67,6 +67,7 @@ struct llama_ubatch {
std::vector<float> embd;
std::vector<float> embd_state;
std::vector<int32_t> embd_row;
+ std::vector<int32_t> embd_state_row; // the state is this row of ctx_state_src's h_nextn device buffer, empty if none
std::vector<llama_pos> pos;
std::vector<int32_t> n_seq_id;
std::vector<llama_seq_id *> seq_id; // these point into the seq_id_data below
@@ -117,6 +118,7 @@ struct llama_batch_ext {
bool has_state = false; // whether state_off is set
size_t state_off = 0; // index offset in the state array
int32_t embd_row = -1; // row of ctx_other's layer input device buffer used as embd, -1 if none
+ int32_t state_row = -1; // row of ctx_state_src's h_nextn device buffer used as state, -1 if none
bool output = false; // TODO: have dedicated output flags
int32_t decision_order = 0; // see llama_batch_ext_set_decision_order()
std::unordered_set<llama_seq_id> seq_ids;
@@ -149,6 +151,7 @@ struct llama_batch_ext {
bool set_token_embd(int32_t idx, llama_embd embd_in);
bool set_token_state(int32_t idx, llama_embd state_in);
bool set_token_embd_row(int32_t idx, int32_t row);
+ bool set_token_state_row(int32_t idx, int32_t row);
bool set_token_pos(int32_t idx, const llama_pos * pos_in);
bool set_output(int32_t idx, bool output_last);
bool set_decision_order(int32_t idx, int32_t order);
@@ -228,6 +231,7 @@ private:
std::vector<int8_t> is_embd_vec; // mixed batch only (= 1 if embd, 0 if text token)
std::vector<float> state_vec; // owned state embeddings built from llama_batch_ext, llama_batch has no slot for them
std::vector<int32_t> embd_row_vec; // owned embd rows built from llama_batch_ext, llama_batch has no slot for them
+ std::vector<int32_t> state_row_vec; // owned state rows built from llama_batch_ext, llama_batch has no slot for them
std::vector<llama_seq_id> seq_id_data; // flat storage for seq_id pointers below
std::vector<llama_pos> pos;
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index a450d8257..394e39d85 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -224,6 +224,8 @@ llama_context::llama_context(
// +1: id n_layer() taps the output of the last layer ("input" of the head)
cparams.embeddings_layer_inp.resize(hparams.n_layer() + 1, false);
cparams.embeddings_layer_inp_dev = false;
+ cparams.embeddings_nextn_dev = false;
+ cparams.ctx_state_src = nullptr;
embd_layer_inp.resize(hparams.n_layer() + 1);
cparams.ctx_type = params.ctx_type;
@@ -1419,6 +1421,12 @@ void llama_context::set_embeddings_nextn(bool value, bool masked) {
cparams.embeddings_nextn = value;
cparams.embeddings_nextn_masked = masked;
+
+ // the device rows are batch rows of the unmasked extraction
+ if (cparams.embeddings_nextn_dev && (!value || masked)) {
+ cparams.embeddings_nextn_dev = false;
+ h_nextn_dev_alloc();
+ }
}
void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) {
@@ -1535,6 +1543,142 @@ void llama_context::layer_inp_dev_alloc() {
}
}
+bool llama_context::h_nextn_dev_alloc() {
+ synchronize();
+
+ auto reset = [&]() {
+ h_nextn_dev = nullptr;
+ h_nextn_dev_dev = nullptr;
+ h_nextn_backend = nullptr;
+ buf_h_nextn_dev.reset();
+ ctx_h_nextn_dev.reset();
+ ev_h_nextn_written_recorded = false;
+ ev_h_nextn_read_recorded = false;
+ };
+
+ // __func__ inside the lambda would be "operator()"
+ const char * func = __func__;
+ auto fail = [&](const char * why) {
+ LLAMA_LOG_WARN("%s: %s, extracting the h_nextn rows to host memory\n", func, why);
+ cparams.embeddings_nextn_dev = false;
+ reset();
+ return false;
+ };
+
+ reset();
+
+ if (!cparams.embeddings_nextn_dev) {
+ return true;
+ }
+
+ h_nextn_dev_dev = model.dev_output();
+ h_nextn_backend = h_nextn_dev_dev ? backend_from_dev(h_nextn_dev_dev) : nullptr;
+ if (h_nextn_backend == nullptr) {
+ return fail("no device backend for the model output");
+ }
+
+ ggml_init_params params = {
+ /*.mem_size =*/ ggml_tensor_overhead(),
+ /*.mem_buffer =*/ nullptr,
+ /*.no_alloc =*/ true,
+ };
+ ctx_h_nextn_dev.reset(ggml_init(params));
+
+ h_nextn_dev = ggml_new_tensor_2d(ctx_h_nextn_dev.get(), GGML_TYPE_F32, model.hparams.n_embd_out(), cparams.n_batch + cparams.n_seq_max);
+ ggml_set_name(h_nextn_dev, "h_nextn_dev");
+
+ buf_h_nextn_dev.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_h_nextn_dev.get(), ggml_backend_dev_buffer_type(h_nextn_dev_dev)));
+ if (!buf_h_nextn_dev) {
+ return fail("failed to allocate the h_nextn device buffer");
+ }
+
+ if (!ev_h_nextn_written) {
+ ev_h_nextn_written.reset(ggml_backend_event_new(h_nextn_dev_dev));
+ ev_h_nextn_read.reset(ggml_backend_event_new(h_nextn_dev_dev));
+ }
+ if (!ev_h_nextn_written || !ev_h_nextn_read) {
+ return fail("the device has no events");
+ }
+
+ LLAMA_LOG_INFO("%s: %10s h_nextn buffer size = %8.2f MiB\n", __func__,
+ ggml_backend_buffer_name(buf_h_nextn_dev.get()), ggml_backend_buffer_get_size(buf_h_nextn_dev.get())/1024.0/1024.0);
+
+ return true;
+}
+
+bool llama_context::set_embeddings_nextn_dev(bool value) {
+ LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
+
+ // the device rows are batch rows, so the extraction must not be masked
+ if (value && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
+ return false;
+ }
+
+ if (cparams.embeddings_nextn_dev != value) {
+ sched_need_reserve = true;
+ }
+
+ cparams.embeddings_nextn_dev = value;
+
+ return h_nextn_dev_alloc() && cparams.embeddings_nextn_dev == value;
+}
+
+bool llama_context::get_h_nextn_rows(int64_t row0, int64_t n_rows, float * out) {
+ if (h_nextn_dev == nullptr || row0 < 0 || n_rows < 0 || row0 + n_rows > h_nextn_dev->ne[1]) {
+ return false;
+ }
+ if (ev_h_nextn_written_recorded) {
+ ggml_backend_event_synchronize(ev_h_nextn_written.get());
+ } else {
+ synchronize();
+ }
+ ggml_backend_tensor_get(h_nextn_dev, out, row0*h_nextn_dev->nb[1], n_rows*h_nextn_dev->nb[1]);
+ return true;
+}
+
+bool llama_context::set_h_nextn_rows(int64_t row0, int64_t n_rows, const float * data) {
+ if (h_nextn_dev == nullptr || row0 < 0 || n_rows < 0 || row0 + n_rows > h_nextn_dev->ne[1]) {
+ return false;
+ }
+ // a reader may still read these rows
+ if (ev_h_nextn_read_recorded) {
+ ggml_backend_event_synchronize(ev_h_nextn_read.get());
+ }
+ ggml_backend_tensor_set(h_nextn_dev, data, row0*h_nextn_dev->nb[1], n_rows*h_nextn_dev->nb[1]);
+ return true;
+}
+
+bool llama_context::set_state_src(llama_context * src) {
+ if (cparams.ctx_state_src != src) {
+ sched_need_reserve = true;
+ }
+ cparams.ctx_state_src = src;
+ return supports_state_row();
+}
+
+bool llama_context::supports_state_row() const {
+ // the MTP graphs read the rows, see llm_graph_context::build_inp_embd_state()
+ if (cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP) {
+ return false;
+ }
+
+ // the rows replace the host state rows, so they must have the width the graph takes for those
+ const ggml_tensor * src = cparams.ctx_state_src ? cparams.ctx_state_src->get_h_nextn_dev() : nullptr;
+ if (src == nullptr || src->ne[0] != (int64_t) llama_batch_ext_select_n_embd_state(cparams.ctx_type, model.hparams)) {
+ return false;
+ }
+
+ // one of our backends must be able to read the buffer
+ ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(src->buffer);
+ for (const auto & backend : backends) {
+ if (ggml_backend_supports_buft(backend.get(), buft)) {
+ return true;
+ }
+ }
+
+ return false;
+}
+
void llama_context::set_nextn_layer_offset(int32_t offset) {
cparams.nextn_layer_offset = offset;
}
@@ -2021,6 +2165,12 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
return -1;
}
+ const bool has_state_row = !batch_inp.tokens.empty() && batch_inp.tokens[0].state_row >= 0;
+ if (has_state_row && !supports_state_row()) {
+ LLAMA_LOG_ERROR("%s: this context cannot read state rows (needs an MTP context and an h_nextn device buffer in ctx_state_src)\n", __func__);
+ return -1;
+ }
+
const auto & vocab = model.vocab;
const auto & hparams = model.hparams;
@@ -2182,6 +2332,22 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
ggml_backend_event_wait(layer_inp_backend, ev_layer_inp_read.get());
}
+ // ordering for the h_nextn device buffer, as for the layer input buffer
+ llama_context * state_src = has_state_row ? cparams.ctx_state_src : nullptr;
+ ggml_backend_t state_backend = nullptr;
+ if (state_src) {
+ state_backend = backend_from_dev(state_src->h_nextn_dev_dev);
+ if (state_src->ev_h_nextn_written_recorded) {
+ ggml_backend_event_synchronize(state_src->ev_h_nextn_written.get());
+ } else {
+ state_src->synchronize();
+ }
+ }
+ const bool write_h_nextn = cparams.embeddings_nextn_dev && h_nextn_dev != nullptr;
+ if (write_h_nextn && ev_h_nextn_read_recorded) {
+ ggml_backend_event_wait(h_nextn_backend, ev_h_nextn_read.get());
+ }
+
do {
const auto & ubatch = mctx->get_ubatch();
@@ -2334,7 +2500,8 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
const int64_t n_rows = masked ? n_outputs : (int64_t) ubatch.n_tokens;
const int64_t offset = masked ? n_outputs_prev : n_tokens_prev;
- if (embd_nextn.data && t_h_nextn && n_rows > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) {
+ // the device buffer replaces the host extraction, see build_h_nextn_dev()
+ if (embd_nextn.data && t_h_nextn && n_rows > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE && !write_h_nextn) {
ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_nextn);
GGML_ASSERT(backend_h != nullptr);
@@ -2371,6 +2538,18 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
if (layer_inp_dev && ev_layer_inp_written) {
ggml_backend_event_record(ev_layer_inp_written.get(), layer_inp_backend);
}
+ if (write_h_nextn) {
+ ggml_backend_event_record(ev_h_nextn_written.get(), h_nextn_backend);
+ ev_h_nextn_written_recorded = true;
+ }
+ if (state_src) {
+ if (state_backend) {
+ ggml_backend_event_record(state_src->ev_h_nextn_read.get(), state_backend);
+ state_src->ev_h_nextn_read_recorded = true;
+ } else {
+ synchronize();
+ }
+ }
if (inject_src) {
if (inject_backend) {
ggml_backend_event_record(inject_src->ev_layer_inp_read.get(), inject_backend);
@@ -2927,6 +3106,8 @@ llm_graph_params llama_context::graph_params(
/*.hdmd =*/ model.hdmd.rot.empty() ? nullptr : &model.hdmd,
/*.layer_inp_dev =*/ layer_inp_dev,
/*.layer_inp_src =*/ supports_embd_row() ? cparams.ctx_other->get_layer_inp_dev() : nullptr,
+ /*.h_nextn_dev =*/ cparams.embeddings_nextn_dev ? h_nextn_dev : nullptr,
+ /*.h_nextn_src =*/ supports_state_row() ? cparams.ctx_state_src->get_h_nextn_dev() : nullptr,
/*.samplers =*/ sampling.samplers,
/*.n_outputs =*/ n_outputs,
/*.cb =*/ graph_get_cb(),
@@ -4530,6 +4711,22 @@ void llama_set_embeddings_nextn(llama_context * ctx, bool value, bool masked) {
ctx->set_embeddings_nextn(value, masked);
}
+bool llama_set_embeddings_nextn_dev(llama_context * ctx, bool value) {
+ return ctx->set_embeddings_nextn_dev(value);
+}
+
+bool llama_get_embeddings_nextn_rows(llama_context * ctx, int64_t row0, int64_t n_rows, float * out) {
+ return ctx->get_h_nextn_rows(row0, n_rows, out);
+}
+
+bool llama_set_embeddings_nextn_rows(llama_context * ctx, int64_t row0, int64_t n_rows, const float * data) {
+ return ctx->set_h_nextn_rows(row0, n_rows, data);
+}
+
+bool llama_set_embd_state_src(llama_context * ctx, llama_context * src) {
+ return ctx->set_state_src(src);
+}
+
void llama_set_embeddings_layer_inp_dev(llama_context * ctx, bool value) {
ctx->set_embeddings_layer_inp_dev(value);
}
diff --git a/src/llama-context.h b/src/llama-context.h
index 39c5619df..ee9c83752 100644
--- a/src/llama-context.h
+++ b/src/llama-context.h
@@ -94,6 +94,7 @@ struct llama_context {
float * get_embeddings_layer_inp(uint32_t lid);
ggml_tensor * get_layer_inp_dev() const { return layer_inp_dev; }
+ ggml_tensor * get_h_nextn_dev() const { return h_nextn_dev; }
llama_token * get_sampled_tokens() const;
llama_token get_sampled_token_ith(int32_t idx);
@@ -126,6 +127,15 @@ struct llama_context {
// whether decode accepts entries whose embd is a row of ctx_other's layer input device buffer
bool supports_embd_row() const;
+
+ // h_nextn device buffer, see llama_set_embeddings_nextn_dev()
+ bool set_embeddings_nextn_dev(bool value);
+ bool get_h_nextn_rows(int64_t row0, int64_t n_rows, float * out);
+ bool set_h_nextn_rows(int64_t row0, int64_t n_rows, const float * data);
+
+ // state rows read from the h_nextn device buffer of src, see llama_set_embd_state_src()
+ bool set_state_src(llama_context * src);
+ bool supports_state_row() const;
void set_nextn_layer_offset(int32_t offset);
void set_causal_attn(bool value);
void set_warmup(bool value);
@@ -252,6 +262,9 @@ private:
// (re)allocate the device buffer for cparams.embeddings_layer_inp_dev
void layer_inp_dev_alloc();
+ // (re)allocate the device buffer for cparams.embeddings_nextn_dev, false if it does not fit
+ bool h_nextn_dev_alloc();
+
ggml_backend_t backend_from_dev(ggml_backend_dev_t dev) const;
//
@@ -445,6 +458,19 @@ private:
ggml_backend_event_ptr ev_layer_inp_read;
bool ev_layer_inp_read_recorded = false;
+ // device buffer with the h_nextn rows, F32 [n_embd_out, n_batch + n_seq_max]:
+ // row r < n_batch is the token at batch index r of the last decode, row n_batch + s is set by the host for seq s
+ // ordered with the reader context like the layer input buffer
+ ggml_context_ptr ctx_h_nextn_dev;
+ ggml_backend_buffer_ptr buf_h_nextn_dev;
+ ggml_tensor * h_nextn_dev = nullptr;
+ ggml_backend_dev_t h_nextn_dev_dev = nullptr;
+ ggml_backend_t h_nextn_backend = nullptr;
+ ggml_backend_event_ptr ev_h_nextn_written;
+ ggml_backend_event_ptr ev_h_nextn_read;
+ bool ev_h_nextn_written_recorded = false;
+ bool ev_h_nextn_read_recorded = false;
+
// keep copies of the per-sequence memory on the device
std::map<llama_seq_id, llama_memory_buffers> mem_storage;
diff --git a/src/llama-cparams.h b/src/llama-cparams.h
index 9d5fe2b3a..7ac1513ef 100644
--- a/src/llama-cparams.h
+++ b/src/llama-cparams.h
@@ -64,6 +64,7 @@ struct llama_cparams {
std::vector<bool> embeddings_layer_inp; // [n_layer()] extract input embeddings for layer
bool embeddings_layer_inp_dev; // write enabled layer inputs to a device buffer instead of host
+ bool embeddings_nextn_dev; // write the unmasked h_nextn rows to a device buffer instead of host, see llama_set_embeddings_nextn_dev()
enum llama_context_type ctx_type;
enum llama_rope_scaling_type rope_scaling_type;
@@ -73,4 +74,7 @@ struct llama_cparams {
void * cb_eval_user_data;
llama_context * ctx_other;
+
+ // context whose h_nextn device buffer holds the state rows of a batch, see llama_set_embd_state_src()
+ llama_context * ctx_state_src;
};
diff --git a/src/llama-ext.h b/src/llama-ext.h
index b55c0c357..a9a08e97e 100644
--- a/src/llama-ext.h
+++ b/src/llama-ext.h
@@ -179,6 +179,24 @@ LLAMA_API int32_t llama_batch_ext_add_embd_row(struct llama_batch_ext * batch, l
LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx);
+// Write the unmasked h_nextn rows (see llama_set_embeddings_nextn) of each decode to a device buffer
+// [n_embd_out, n_batch + n_seq_max] instead of host memory: row r < n_batch is the token at batch index r of the last decode,
+// row n_batch + s is free for the caller, e.g. for seq s (see llama_set_embeddings_nextn_rows()).
+// Returns false if the buffer cannot be used, the rows are then extracted to host memory.
+LLAMA_API bool llama_set_embeddings_nextn_dev(struct llama_context * ctx, bool value);
+
+// Copy rows of the h_nextn device buffer to or from host memory, ordered after the last decode and its readers
+LLAMA_API bool llama_get_embeddings_nextn_rows(struct llama_context * ctx, int64_t row0, int64_t n_rows, float * out);
+LLAMA_API bool llama_set_embeddings_nextn_rows(struct llama_context * ctx, int64_t row0, int64_t n_rows, const float * data);
+
+// Read the state rows of batches (see llama_batch_ext_set_embd_state_row()) from the h_nextn device buffer of src.
+// Returns whether decode accepts batches with state rows.
+LLAMA_API bool llama_set_embd_state_src(struct llama_context * ctx, struct llama_context * src);
+
+// The state of entry idx is row `row` of the h_nextn device buffer of the context set with llama_set_embd_state_src().
+// All entries of a batch have a state row, or none.
+LLAMA_API bool llama_batch_ext_set_embd_state_row(struct llama_batch_ext * batch, int32_t idx, int32_t row);
+
//
// model/context data extraction
//
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index b38daca76..0f7d8ab17 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -223,6 +223,12 @@ void llm_graph_input_embd_h::set_input(const llama_ubatch * ubatch) {
ggml_backend_tensor_set(embd, ubatch->embd, 0, n_tokens*n_embd*ggml_element_size(embd));
}
+ if (h_rows) {
+ GGML_ASSERT(ubatch->data && ubatch->data->embd_state_row.size() == (size_t) n_tokens);
+ ggml_backend_tensor_set(h_rows, ubatch->data->embd_state_row.data(), 0, n_tokens*ggml_element_size(h_rows));
+ return;
+ }
+
GGML_ASSERT(ubatch->embd_state && "this graph requires a state embedding, see llama_batch_ext_set_embd_state()");
GGML_ASSERT(n_embd_state == h->ne[0]);
@@ -236,6 +242,10 @@ bool llm_graph_input_embd_h::can_reuse(const llm_graph_params & params) {
res &= (!params.ubatch.embd) || (embd && embd->ne[1] == params.ubatch.n_tokens);
res &= (!params.ubatch.embd_state) || (h && h->ne[1] == params.ubatch.n_tokens);
+ const bool has_rows = params.ubatch.data && !params.ubatch.data->embd_state_row.empty();
+ res &= has_rows == (h_rows != nullptr);
+ res &= !has_rows || h_rows->ne[0] == params.ubatch.n_tokens;
+
return res;
}
@@ -1588,6 +1598,8 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
hdmd (params.hdmd),
layer_inp_dev (params.layer_inp_dev),
layer_inp_src (params.layer_inp_src),
+ h_nextn_dev (params.h_nextn_dev),
+ h_nextn_src (params.h_nextn_src),
samplers (params.samplers),
cb_func (params.cb),
res (params.res),
@@ -2883,6 +2895,24 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float to
return cur;
}
+ggml_tensor * llm_graph_context::build_inp_embd_state(llm_graph_input_embd_h * inp, int64_t n_embd_state) const {
+ if (ubatch.data && !ubatch.data->embd_state_row.empty()) {
+ GGML_ASSERT(h_nextn_src != nullptr && h_nextn_src->ne[0] == n_embd_state);
+
+ inp->h_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp->h_rows);
+ ggml_set_name(inp->h_rows, "mtp_h_rows");
+
+ return ggml_get_rows(ctx0, h_nextn_src, inp->h_rows);
+ }
+
+ ggml_tensor * h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_state, n_tokens);
+ ggml_set_input(h);
+ ggml_set_name(h, "mtp_h_input");
+
+ return h;
+}
+
ggml_tensor * llm_graph_context::build_inp_pos() const {
auto inp = std::make_unique<llm_graph_input_pos>(hparams.n_pos_per_embd(), hparams.rope_section_order);
@@ -4396,6 +4426,25 @@ void llm_graph_context::build_layer_inp_dev() const {
res->add_input(std::move(inp));
}
+void llm_graph_context::build_h_nextn_dev() const {
+ ggml_tensor * src = res->t_h_nextn;
+ if (h_nextn_dev == nullptr || src == nullptr) {
+ return;
+ }
+
+ // the device buffer is only enabled for the unmasked extraction, which has a row per batch token
+ GGML_ASSERT(src->ne[0] == h_nextn_dev->ne[0] && ggml_nrows(src) == n_tokens);
+
+ auto inp = std::make_unique<llm_graph_input_batch_idxs>();
+
+ inp->idxs = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_tokens);
+ ggml_set_input(inp->idxs);
+
+ ggml_build_forward_expand(gf, ggml_set_rows(ctx0, h_nextn_dev, ggml_reshape_2d(ctx0, src, src->ne[0], n_tokens), inp->idxs));
+
+ res->add_input(std::move(inp));
+}
+
void llm_graph_context::build_sampling() const {
if (samplers.empty() || !res->t_logits) {
return;
diff --git a/src/llama-graph.h b/src/llama-graph.h
index 9867c9540..20c43e374 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -167,7 +167,8 @@ public:
ggml_tensor * tokens = nullptr; // I32 [n_batch]
ggml_tensor * embd = nullptr; // F32 [n_embd, n_batch]
- ggml_tensor * h = nullptr; // F32 [n_embd_state, n_batch]
+ ggml_tensor * h = nullptr; // F32 [n_embd_state, n_batch], the host input or the rows gathered with h_rows
+ ggml_tensor * h_rows = nullptr; // I32 [n_batch], rows of the h_nextn device buffer of ctx_state_src, nullptr on the host path
const int64_t n_embd = 0;
const int64_t n_embd_state = 0;
@@ -812,6 +813,8 @@ struct llm_graph_params {
ggml_tensor * layer_inp_dev = nullptr; // own layer input device buffer (write)
ggml_tensor * layer_inp_src = nullptr; // ctx_other's layer input device buffer (read)
+ ggml_tensor * h_nextn_dev = nullptr; // own h_nextn device buffer (write)
+ ggml_tensor * h_nextn_src = nullptr; // ctx_state_src's h_nextn device buffer (read)
std::map<llama_seq_id, llama_sampler *> samplers;
@@ -907,6 +910,8 @@ struct llm_graph_params {
cparams.causal_attn == other.cparams.causal_attn &&
layer_inp_dev == other.layer_inp_dev &&
layer_inp_src == other.layer_inp_src &&
+ h_nextn_dev == other.h_nextn_dev &&
+ h_nextn_src == other.h_nextn_src &&
model == other.model &&
gtype == other.gtype &&
cvec == other.cvec &&
@@ -1081,6 +1086,8 @@ struct llm_graph_context {
ggml_tensor * layer_inp_dev;
ggml_tensor * layer_inp_src;
+ ggml_tensor * h_nextn_dev;
+ ggml_tensor * h_nextn_src;
std::map<llama_seq_id, llama_sampler *> samplers;
@@ -1256,6 +1263,11 @@ struct llm_graph_context {
// tok_scale: applied to token rows only
ggml_tensor * build_inp_embd(ggml_tensor * tok_embd, float tok_scale = 1.0f) const;
+
+ // the state input of an MTP graph: rows of the h_nextn device buffer of ctx_state_src (see llama_set_embd_state_src()),
+ // or uploaded from ubatch.embd_state
+ ggml_tensor * build_inp_embd_state(llm_graph_input_embd_h * inp, int64_t n_embd_state) const;
+
ggml_tensor * build_inp_pos() const;
// applies the RoPE of layer il, cur is returned as is if the layer has none
@@ -1507,6 +1519,9 @@ struct llm_graph_context {
// write enabled layer inputs into layer_inp_dev at their batch rows
void build_layer_inp_dev() const;
+ // write the unmasked h_nextn rows into h_nextn_dev at their batch rows
+ void build_h_nextn_dev() const;
+
//
// dense (out)
//
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 7e845bc30..72639b155 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -3516,6 +3516,8 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
llm->build_layer_inp_dev();
+ llm->build_h_nextn_dev();
+
llm->res->set_outputs(params);
return llm->res->get_gf();
diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp
index 8e47e2a0a..ad1f0ee27 100644
--- a/src/models/bailingmoe3.cpp
+++ b/src/models/bailingmoe3.cpp
@@ -404,9 +404,7 @@ llama_model_bailingmoe3::graph_mtp::graph_mtp(const llm_graph_params & params) :
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * tok_embd = ubatch.token ? ggml_get_rows(ctx0, model.tok_embd, inp->tokens) : inp->embd;
ggml_tensor * h_norm = build_norm(inp->h, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
diff --git a/src/models/cohere2moe.cpp b/src/models/cohere2moe.cpp
index 5bb2db15a..52b959672 100644
--- a/src/models/cohere2moe.cpp
+++ b/src/models/cohere2moe.cpp
@@ -316,9 +316,7 @@ llama_model_cohere2moe::graph_mtp::graph_mtp(const llm_graph_params & params) :
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp
index 7ed35ee9a..96c628ff7 100644
--- a/src/models/deepseek2.cpp
+++ b/src/models/deepseek2.cpp
@@ -227,9 +227,7 @@ llama_model_deepseek2::graph_mtp::graph_mtp(const llm_graph_params & params) :
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp
index 6a7a0a7a9..fc6207f8b 100644
--- a/src/models/deepseek32.cpp
+++ b/src/models/deepseek32.cpp
@@ -531,9 +531,7 @@ llama_model_deepseek32::graph_mtp::graph_mtp(const llm_graph_params & params)
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp
index 9345fa833..0d5d9a469 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -1386,9 +1386,7 @@ llama_model_deepseek4::graph_mtp::graph_mtp(const llm_graph_params & params) :
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd_out());
ggml_tensor * tok_embd;
if (ubatch.token) {
diff --git a/src/models/gemma4-assistant.cpp b/src/models/gemma4-assistant.cpp
index 77968cbb3..48b7aa864 100644
--- a/src/models/gemma4-assistant.cpp
+++ b/src/models/gemma4-assistant.cpp
@@ -103,9 +103,8 @@ llama_model_gemma4_assistant::graph::graph(const llm_graph_params & params) :
inp_embd = inp->embd;
res->t_inp_embd = inp->embd;
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_backbone, ubatch.n_tokens);
+ inp->h = build_inp_embd_state(inp.get(), n_embd_backbone);
cb(inp->h, "inp_h", -1);
- ggml_set_input(inp->h);
inp_h = inp->h;
res->add_input(std::move(inp));
diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp
index 358defe83..7007f427f 100644
--- a/src/models/glm-dsa.cpp
+++ b/src/models/glm-dsa.cpp
@@ -579,9 +579,7 @@ llama_model_glm_dsa::graph_mtp::graph_mtp(const llm_graph_params & params)
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/glm4-moe.cpp b/src/models/glm4-moe.cpp
index 2d958aa95..9560fd3d0 100644
--- a/src/models/glm4-moe.cpp
+++ b/src/models/glm4-moe.cpp
@@ -161,9 +161,7 @@ llama_model_glm4_moe::graph_mtp::graph_mtp(const llm_graph_params & params)
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index 5789b3013..844a327b3 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -564,9 +564,7 @@ llama_model_glm5_next::graph_mtp::graph_mtp(const llm_graph_params & params) :
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), n_embd);
ggml_tensor * tok_embd;
if (ubatch.token) {
diff --git a/src/models/hy-v3.cpp b/src/models/hy-v3.cpp
index 3fce6a246..c47c373f8 100644
--- a/src/models/hy-v3.cpp
+++ b/src/models/hy-v3.cpp
@@ -246,9 +246,7 @@ llama_model_hy_v3::graph_mtp::graph_mtp(const llm_graph_params & params)
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
diff --git a/src/models/mellum.cpp b/src/models/mellum.cpp
index 83f70019d..57dd5b79a 100644
--- a/src/models/mellum.cpp
+++ b/src/models/mellum.cpp
@@ -266,9 +266,7 @@ llama_model_mellum::graph_mtp<iswa>::graph_mtp(const llm_graph_params & params)
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
// TODO: make static using `ggml_build_forward_select()`
// see llm_graph_context::build_inp_embd() for reference
diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp
index b318b05ed..4fa97ff5b 100644
--- a/src/models/mimo2.cpp
+++ b/src/models/mimo2.cpp
@@ -275,9 +275,7 @@ llama_model_mimo2::graph_mtp::graph_mtp(const llm_graph_params & params)
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
ggml_tensor * h_input = inp->h;
diff --git a/src/models/nemotron-h-moe.cpp b/src/models/nemotron-h-moe.cpp
index 871b2d961..475355832 100644
--- a/src/models/nemotron-h-moe.cpp
+++ b/src/models/nemotron-h-moe.cpp
@@ -42,9 +42,7 @@ llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llm_graph_params & params
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp
index d0b4f4c2c..736e23f10 100644
--- a/src/models/qwen35.cpp
+++ b/src/models/qwen35.cpp
@@ -528,9 +528,7 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llm_graph_params & params)
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp
index 0b06b1f5d..c4da3e2be 100644
--- a/src/models/qwen35moe.cpp
+++ b/src/models/qwen35moe.cpp
@@ -567,9 +567,7 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llm_graph_params & params)
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp
index 0731838f8..110581bae 100644
--- a/src/models/qwen3next.cpp
+++ b/src/models/qwen3next.cpp
@@ -647,9 +647,7 @@ llama_model_qwen3next::graph_mtp::graph_mtp(const llm_graph_params & params)
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index c543edf0a..876640e16 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -596,9 +596,7 @@ llama_model_qwen4exp::graph_mtp::graph_mtp(const llm_graph_params & params) :
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd_out());
ggml_tensor * tok_embd;
if (ubatch.token) {
diff --git a/src/models/step35.cpp b/src/models/step35.cpp
index ad8a3a1c7..5bed70163 100644
--- a/src/models/step35.cpp
+++ b/src/models/step35.cpp
@@ -379,9 +379,7 @@ llama_model_step35::graph_mtp::graph_mtp(const llm_graph_params & params)
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
diff --git a/src/models/xing4-0.cpp b/src/models/xing4-0.cpp
index c95acd6b8..f405fcc9c 100644
--- a/src/models/xing4-0.cpp
+++ b/src/models/xing4-0.cpp
@@ -237,9 +237,7 @@ llama_model_xing4_0::graph_mtp::graph_mtp(const llama_model & model, const llm_g
}
cb(tok_embd, "mtp_tok_embd", il);
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
- ggml_set_input(inp->h);
- ggml_set_name(inp->h, "mtp_h_input");
+ inp->h = build_inp_embd_state(inp.get(), hparams.n_embd);
ggml_tensor * h_embd = inp->h;