Commit 2cbca389f for llama.cpp
commit 2cbca389f515df455412dafa07e2b1359cc72a6d
Author: Aman Gupta <amangupta052@gmail.com>
Date: Tue Nov 3 09:02:32 2026 +0800
dflash: copy hidden features via device copies (#31506)
diff --git a/common/common.cpp b/common/common.cpp
index 04e348805..95a2dceef 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -2275,6 +2275,16 @@ int32_t common_batch::add_embd(llama_embd embd, const llama_pos * pos, llama_seq
return size() - 1;
}
+int32_t common_batch::add_embd_row(int32_t row, const llama_pos * pos, llama_seq_id seq_id, bool output) {
+ token t = { LLAMA_TOKEN_NULL, { 0, 0, 0, 0 }, seq_id, output, { nullptr, 0, 0 }, { nullptr, 0, 0 }, {} };
+ t.embd_row = row;
+ for (int32_t j = 0; j < n_pos; ++j) {
+ t.pos[j] = pos[j];
+ }
+ tokens.push_back(t);
+ return size() - 1;
+}
+
llama_batch_ext * common_batch::get_sub_batch(int32_t off, int32_t n) {
GGML_ASSERT(batch && "common_batch was not initialized with a context");
GGML_ASSERT(off >= 0 && n >= 0 && off + n <= size());
@@ -2295,6 +2305,12 @@ llama_batch_ext * common_batch::get_sub_batch(int32_t off, int32_t n) {
if (t.embd.data && !llama_batch_ext_set_embd_token(res, idx, t.embd)) {
GGML_ABORT("%s: failed to set the embedding of token %d at index %d\n", __func__, t.id, i);
}
+ } else if (t.embd_row >= 0) {
+ idx = llama_batch_ext_add_embd_row(res, t.seq_id, t.embd_row);
+ if (idx < 0) {
+ GGML_ABORT("%s: failed to add embd row %d at index %d (error %d, n = %d)\n", __func__, t.embd_row, i, idx, n);
+ }
+ llama_batch_ext_set_pos(res, idx, t.pos.data());
} else {
idx = llama_batch_ext_add_embd(res, t.seq_id, t.embd);
if (idx < 0) {
diff --git a/common/common.h b/common/common.h
index d0e33fb75..e970fb39b 100644
--- a/common/common.h
+++ b/common/common.h
@@ -1096,6 +1096,7 @@ struct common_batch {
llama_embd state; // non-owning view of the data passed to set_embd_state(), data == NULL if none
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
};
std::vector<token> tokens; // mirror of the entries, tokens[i] describes batch index i
@@ -1138,6 +1139,10 @@ struct common_batch {
// pos points to n_pos positions
int32_t add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output);
+ // add an entry whose embedding is row `row` of the layer input device buffer of the context's ctx_other
+ // pos points to n_pos positions
+ int32_t add_embd_row(int32_t row, const llama_pos * pos, llama_seq_id seq_id, bool output);
+
int32_t size() const { return (int32_t) tokens.size(); }
};
diff --git a/common/speculative.cpp b/common/speculative.cpp
index 4305ac6ac..14fee2b2c 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -1027,6 +1027,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices
uint32_t target_layer_ids_n = 0;
+ // inject target features from the target's device buffer instead of host copies (env LLAMA_DFLASH_DEV=0 disables)
+ bool use_dev_features = false;
+
common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,
common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
: common_speculative_impl(type, n_seq, params.draft.n_max)
@@ -1141,6 +1144,23 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
}
+ // the device buffer keeps layers in ascending id order, the fc input expects target_layer_ids order
+ {
+ const char * env = getenv("LLAMA_DFLASH_DEV");
+ const bool ascending = std::adjacent_find(target_layer_ids, target_layer_ids + target_layer_ids_n,
+ [](int32_t a, int32_t b) { return a >= b; }) == target_layer_ids + target_layer_ids_n;
+
+ use_dev_features = !is_lilicorr && ascending && (env == nullptr || atoi(env) != 0);
+ if (use_dev_features) {
+ llama_set_embeddings_layer_inp_dev(ctx_tgt, true);
+ use_dev_features = llama_supports_embd_row(ctx_dft);
+ if (!use_dev_features) {
+ llama_set_embeddings_layer_inp_dev(ctx_tgt, false);
+ }
+ }
+ LOG_INF("%s: DFlash target features via %s\n", __func__, use_dev_features ? "device buffer" : "host copies");
+ }
+
// DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2 && !is_lilicorr);
llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise
@@ -1233,6 +1253,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
+ if (use_dev_features) {
+ // the features are the target batch rows in the target's device buffer
+ batch_inject.clear();
+ for (int32_t i = 0; i < n_chunk; ++i) {
+ const int32_t row = i_batch_beg[seq_id] + offset + i;
+ const llama_pos p = batch_in.tokens[row].pos[0];
+ const llama_pos pos_arr[4] = { p, p, p, 0 };
+ batch_inject.add_embd_row(row, pos_arr, seq_id, false);
+ }
+
+ const int32_t rc = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch_inject.get());
+ if (rc != 0) {
+ LOG_ERR("%s: llama_process(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+ __func__, rc, (int) n_chunk, (int) offset);
+ return false;
+ }
+ continue;
+ }
+
// gather target features per extract layer; the fused decode encodes and
// injects them into the K/V cache at the target positions
features_buf.resize((size_t) n_chunk * n_embd_enc);
diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 321aa9d23..f84c6d50b 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -52,6 +52,7 @@ bool llama_batch_allocr::init(
// an entry can carry a token id, a token embedding, or both
// all entries must carry the same combination, or be a mix of token and embd entries
// a state embedding (e.g. MTP hook batches) is set on all entries or on none
+ // an embd row (embd read from a device buffer) is set on all entries or on none, with no token id or embedding
//
int32_t n_tok_only = 0;
@@ -59,12 +60,24 @@ bool llama_batch_allocr::init(
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;
for (int32_t i = 0; i < n_tok; ++i) {
const bool is_tok = batch_inp.tokens[i].id != LLAMA_TOKEN_NULL;
const bool is_emb = batch_inp.tokens[i].has_embd;
+ const bool is_row = batch_inp.tokens[i].embd_row >= 0;
- if (!is_tok && !is_emb) {
+ if (is_row != has_row) {
+ LLAMA_LOG_ERROR("%s: all entries in the batch must have the same embd row presence\n", __func__);
+ return false;
+ }
+
+ if (is_row) {
+ if (is_tok || is_emb) {
+ LLAMA_LOG_ERROR("%s: entry %d has an embd row and a token id or embedding\n", __func__, i);
+ return false;
+ }
+ } else if (!is_tok && !is_emb) {
LLAMA_LOG_ERROR("%s: entry %d has neither a token id nor an embedding\n", __func__, i);
return false;
}
@@ -137,6 +150,13 @@ bool llama_batch_allocr::init(
state_vec = batch_inp.state;
}
+ if (has_row) {
+ embd_row_vec.resize(n_tok);
+ for (int32_t i = 0; i < n_tok; ++i) {
+ embd_row_vec[i] = batch_inp.tokens[i].embd_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)
@@ -306,6 +326,7 @@ bool llama_batch_allocr::init(
/*.token =*/ batch.token,
/*.embd =*/ batch.embd,
/*.embd_state =*/ state_vec.empty() ? nullptr : state_vec.data(),
+ /*.embd_row =*/ embd_row_vec.empty() ? nullptr : embd_row_vec.data(),
/*.pos =*/ batch.pos,
/*.n_seq_id =*/ batch.n_seq_id,
/*.seq_id =*/ batch.seq_id,
@@ -531,6 +552,7 @@ llama_ubatch llama_batch_allocr::ubatch_reserve(uint32_t n_seq_tokens, uint32_t
/*.token =*/ udata->token.data(),
/*.embd =*/ nullptr,
/*.embd_state =*/ nullptr,
+ /*.embd_row =*/ nullptr,
/*.pos =*/ udata->pos.data(),
/*.n_seq_id =*/ udata->n_seq_id.data(),
/*.seq_id =*/ udata->seq_id.data(),
@@ -838,6 +860,7 @@ void llama_batch_allocr::clear() {
embd_vec .clear();
is_embd_vec .clear();
state_vec .clear();
+ embd_row_vec.clear();
seq_id_data .clear();
pos .clear();
n_seq_id .clear();
@@ -881,6 +904,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
const bool use_token = batch.token && !(mixed_batch && n_embd_rows == n_tokens);
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 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;
@@ -889,6 +913,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
udata->token .resize(n_tokens);
udata->embd .resize(n_embd_all);
udata->embd_state.resize(n_state_all);
+ udata->embd_row .resize(has_row ? n_tokens : 0);
udata->pos .resize(n_pos_all);
udata->n_seq_id .resize(n_tokens);
udata->seq_id .resize(n_tokens);
@@ -920,6 +945,10 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
memcpy(udata->embd_state.data() + i*n_embd_state, state_vec.data() + (int64_t) idxs[i]*n_embd_state, n_embd_state*sizeof(float));
}
+ if (has_row) {
+ udata->embd_row[i] = embd_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]];
}
@@ -967,6 +996,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector<int32_t> & idxs, u
/*.token =*/ use_token ? udata->token.data() : nullptr,
/*.embd =*/ use_embd ? udata->embd.data() : nullptr,
/*.embd_state =*/ has_state ? udata->embd_state.data() : nullptr,
+ /*.embd_row =*/ has_row ? udata->embd_row.data() : nullptr,
/*.pos =*/ udata->pos.data(),
/*.n_seq_id =*/ udata->n_seq_id.data(),
/*.seq_id =*/ udata->seq_id.data(),
@@ -1019,6 +1049,7 @@ void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) {
LLAMA_LOG_DEBUG("%s: token = %p\n", __func__, (void *) ubatch.token);
LLAMA_LOG_DEBUG("%s: embd = %p\n", __func__, (void *) ubatch.embd);
LLAMA_LOG_DEBUG("%s: embd_state = %p\n", __func__, (void *) ubatch.embd_state);
+ LLAMA_LOG_DEBUG("%s: embd_row = %p\n", __func__, (void *) ubatch.embd_row);
LLAMA_LOG_DEBUG("%s: pos = %p\n", __func__, (void *) ubatch.pos);
LLAMA_LOG_DEBUG("%s: n_seq_id = %p\n", __func__, (void *) ubatch.n_seq_id);
LLAMA_LOG_DEBUG("%s: seq_id = %p\n", __func__, (void *) ubatch.seq_id);
@@ -1303,6 +1334,26 @@ bool llama_batch_ext::set_token_state(int32_t idx, llama_embd state_in) {
return true;
}
+bool llama_batch_ext::set_token_embd_row(int32_t idx, int32_t row) {
+ if (idx < 0 || idx >= (int32_t) tokens.size()) {
+ return false;
+ }
+ if (row < 0) {
+ return false;
+ }
+
+ token & t = tokens[idx];
+
+ if (t.id != LLAMA_TOKEN_NULL || t.has_embd || t.embd_row >= 0) {
+ LLAMA_LOG_ERROR("%s: entry %d already has a token id, an embedding or an embd row\n", __func__, idx);
+ return false;
+ }
+
+ t.embd_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;
@@ -1377,6 +1428,17 @@ int32_t llama_batch_ext_add_embd(llama_batch_ext * batch, llama_seq_id seq_id, l
return idx;
}
+int32_t llama_batch_ext_add_embd_row(llama_batch_ext * batch, llama_seq_id seq_id, int32_t row) {
+ int32_t idx = batch->add_token(seq_id);
+ if (idx < 0) {
+ return idx;
+ }
+ if (!batch->set_token_embd_row(idx, row)) {
+ return -2;
+ }
+ return idx;
+}
+
bool llama_batch_ext_add_seq(llama_batch_ext * batch, int32_t idx, llama_seq_id seq_id) {
return batch->add_seq(idx, seq_id);
}
diff --git a/src/llama-batch.h b/src/llama-batch.h
index e1e43979b..d2361396c 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -52,6 +52,7 @@ struct llama_ubatch {
llama_token * token; // [n_tokens] | i | id, token
float * embd; // [n_embd, n_tokens] | i | embd
float * embd_state; // [n_embd_state, n_tokens] | i | hidden state carried over from a previous stage (e.g. MTP)
+ int32_t * embd_row; // [n_tokens] | i | embd is this row of ctx_other's layer input device buffer
llama_pos * pos; // [n_tokens*n_pos] | i | pos
int32_t * n_seq_id; // [n_tokens] | i | -
llama_seq_id ** seq_id; // [n_tokens] | s | s0, s1, seq_id
@@ -65,6 +66,7 @@ struct llama_ubatch {
std::vector<llama_token> token;
std::vector<float> embd;
std::vector<float> embd_state;
+ std::vector<int32_t> embd_row;
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
@@ -114,6 +116,7 @@ struct llama_batch_ext {
size_t embd_off = 0; // index offset in the embd array
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
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;
@@ -145,6 +148,7 @@ struct llama_batch_ext {
bool set_token_id(int32_t idx, llama_token id);
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_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);
@@ -223,6 +227,7 @@ private:
std::vector<float> embd_vec; // owned embeddings built from llama_batch_ext
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<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 66b724c18..0cd795a56 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -222,6 +222,7 @@ 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;
embd_layer_inp.resize(hparams.n_layer() + 1);
cparams.ctx_type = params.ctx_type;
@@ -257,6 +258,9 @@ llama_context::llama_context(
throw std::runtime_error(model.arch_name() + " requires ctx_other to be set (this warning is normal during memory fitting)");
}
cparams.ctx_other = params.ctx_other;
+ } else if (model.arch == LLM_ARCH_DFLASH) {
+ // used to read the target's layer input device buffer
+ cparams.ctx_other = params.ctx_other;
}
}
@@ -1396,6 +1400,109 @@ void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) {
// note: without this reserve, the draft acceptance drops to zero. not sure why - this is unexpected
sched_need_reserve = true;
+
+ if (cparams.embeddings_layer_inp_dev) {
+ layer_inp_dev_alloc();
+ }
+}
+
+void llama_context::set_embeddings_layer_inp_dev(bool value) {
+ LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
+
+ if (cparams.embeddings_layer_inp_dev != value) {
+ sched_need_reserve = true;
+ }
+
+ cparams.embeddings_layer_inp_dev = value;
+
+ layer_inp_dev_alloc();
+}
+
+bool llama_context::supports_embd_row() const {
+ // the DFlash/DSpark decoder graphs read the device buffer
+ if (model.arch != LLM_ARCH_DFLASH) {
+ return false;
+ }
+
+ const ggml_tensor * src = cparams.ctx_other ? cparams.ctx_other->get_layer_inp_dev() : nullptr;
+ if (src == nullptr) {
+ return false;
+ }
+
+ // one of our backends must be able to read the buffer, e.g. not a draft on another GPU than the target output
+ 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;
+}
+
+ggml_backend_t llama_context::backend_from_dev(ggml_backend_dev_t dev) const {
+ for (const auto & backend : backends) {
+ if (ggml_backend_get_device(backend.get()) == dev) {
+ return backend.get();
+ }
+ }
+ return nullptr;
+}
+
+void llama_context::layer_inp_dev_alloc() {
+ synchronize();
+
+ layer_inp_dev = nullptr;
+ layer_inp_dev_dev = nullptr;
+ layer_inp_backend = nullptr;
+ buf_layer_inp_dev.reset();
+ ctx_layer_inp_dev.reset();
+
+ if (!cparams.embeddings_layer_inp_dev) {
+ return;
+ }
+
+ int64_t n_layer_inp = 0;
+ for (bool enabled : cparams.embeddings_layer_inp) {
+ n_layer_inp += enabled ? 1 : 0;
+ }
+ if (n_layer_inp == 0) {
+ return;
+ }
+
+ ggml_init_params params = {
+ /*.mem_size =*/ ggml_tensor_overhead(),
+ /*.mem_buffer =*/ nullptr,
+ /*.no_alloc =*/ true,
+ };
+ ctx_layer_inp_dev.reset(ggml_init(params));
+
+ layer_inp_dev = ggml_new_tensor_2d(ctx_layer_inp_dev.get(), GGML_TYPE_F32, n_layer_inp*model.hparams.n_embd, cparams.n_batch);
+ ggml_set_name(layer_inp_dev, "layer_inp_dev");
+
+ layer_inp_dev_dev = model.dev_output();
+ auto * buft = layer_inp_dev_dev ? ggml_backend_dev_buffer_type(layer_inp_dev_dev) : ggml_backend_cpu_buffer_type();
+
+ buf_layer_inp_dev.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_layer_inp_dev.get(), buft));
+ if (!buf_layer_inp_dev) {
+ // e.g. a large n_batch that does not fit next to the model, keep the host copies
+ LLAMA_LOG_WARN("%s: failed to allocate the layer input device buffer, using host copies\n", __func__);
+ cparams.embeddings_layer_inp_dev = false;
+ layer_inp_dev = nullptr;
+ layer_inp_dev_dev = nullptr;
+ ctx_layer_inp_dev.reset();
+ return;
+ }
+
+ LLAMA_LOG_INFO("%s: %10s layer input buffer size = %8.2f MiB (%d layers)\n", __func__,
+ ggml_backend_buffer_name(buf_layer_inp_dev.get()), ggml_backend_buffer_get_size(buf_layer_inp_dev.get())/1024.0/1024.0, (int) n_layer_inp);
+
+ layer_inp_backend = layer_inp_dev_dev ? backend_from_dev(layer_inp_dev_dev) : nullptr;
+ if (layer_inp_backend && !ev_layer_inp_written) {
+ ev_layer_inp_written.reset(ggml_backend_event_new(layer_inp_dev_dev));
+ ev_layer_inp_read.reset(ggml_backend_event_new(layer_inp_dev_dev));
+ ev_layer_inp_read_recorded = false;
+ }
}
void llama_context::set_nextn_layer_offset(int32_t offset) {
@@ -1627,6 +1734,11 @@ int llama_context::encode(const llama_batch_ext & batch_inp) {
return -1;
}
+ if (!batch_inp.tokens.empty() && batch_inp.tokens[0].embd_row >= 0) {
+ LLAMA_LOG_ERROR("%s: embd row batches are not supported by encode\n", __func__);
+ return -1;
+ }
+
// eagle3/DFlash: features as encoder input, and non-draft paths fall back to model's input dim
const int64_t n_vocab = model.vocab.n_tokens();
@@ -1873,6 +1985,12 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
return -1;
}
+ const bool has_embd_row = !batch_inp.tokens.empty() && batch_inp.tokens[0].embd_row >= 0;
+ if (has_embd_row && !supports_embd_row()) {
+ LLAMA_LOG_ERROR("%s: this context cannot read embd rows (needs a DFlash/DSpark decoder and a layer input device buffer in ctx_other)\n", __func__);
+ return -1;
+ }
+
const auto & vocab = model.vocab;
const auto & hparams = model.hparams;
@@ -2016,6 +2134,24 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
int64_t n_outputs_prev = 0;
int64_t n_tokens_prev = 0;
+ // ordering for the layer input device buffer
+ llama_context * inject_src = has_embd_row ? cparams.ctx_other : nullptr;
+ ggml_backend_t inject_backend = nullptr;
+ if (inject_src) {
+ inject_backend = inject_src->ev_layer_inp_written ? backend_from_dev(inject_src->layer_inp_dev_dev) : nullptr;
+ if (inject_backend) {
+ // wait on the host: a device-side wait would stall our queue for the whole target batch,
+ // which trips the Metal command buffer timeout on long prompts
+ ggml_backend_event_synchronize(inject_src->ev_layer_inp_written.get());
+ } else {
+ inject_src->synchronize();
+ }
+ }
+ if (layer_inp_dev && ev_layer_inp_read && ev_layer_inp_read_recorded) {
+ // do not overwrite rows the reader has not consumed yet
+ ggml_backend_event_wait(layer_inp_backend, ev_layer_inp_read.get());
+ }
+
do {
const auto & ubatch = mctx->get_ubatch();
@@ -2202,6 +2338,18 @@ int llama_context::decode(const llama_batch_ext & batch_inp) {
n_tokens_prev += ubatch.n_tokens;
} while (mctx->next());
+ if (layer_inp_dev && ev_layer_inp_written) {
+ ggml_backend_event_record(ev_layer_inp_written.get(), layer_inp_backend);
+ }
+ if (inject_src) {
+ if (inject_backend) {
+ ggml_backend_event_record(inject_src->ev_layer_inp_read.get(), inject_backend);
+ inject_src->ev_layer_inp_read_recorded = true;
+ } else {
+ synchronize();
+ }
+ }
+
// set to total number of outputs in the batch, for use in llama_get_logits_ith
n_outputs = n_outputs_all;
@@ -2298,7 +2446,7 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) {
}
for (bool enabled : cparams.embeddings_layer_inp) {
- if (enabled) {
+ if (enabled && !cparams.embeddings_layer_inp_dev) {
embd_layer_inp_float_count += (size_t) n_embd * n_batch;
}
}
@@ -2370,7 +2518,7 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) {
offset += embd_nextn.size * sizeof(float);
for (uint32_t il = 0; il < embd_layer_inp.size(); ++il) {
- if (cparams.embeddings_layer_inp[il]) {
+ if (cparams.embeddings_layer_inp[il] && !cparams.embeddings_layer_inp_dev) {
embd_layer_inp[il] = buffer_view<float>{(float *) (base + offset), (size_t) n_embd * n_batch};
offset += embd_layer_inp[il].size * sizeof(float);
} else {
@@ -2426,6 +2574,9 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) {
bool llama_context::extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens) {
bool extracted = false;
+ if (cparams.embeddings_layer_inp_dev) {
+ return extracted;
+ }
for (uint32_t il = 0; il < cparams.embeddings_layer_inp.size(); ++il) {
if (!cparams.embeddings_layer_inp[il]) {
continue;
@@ -2742,6 +2893,8 @@ llm_graph_params llama_context::graph_params(
/*.moe_cache =*/ moe_cache.get(),
/*.prec_policy =*/ &model.prec_policy,
/*.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,
/*.samplers =*/ sampling.samplers,
/*.n_outputs =*/ n_outputs,
/*.cb =*/ graph_get_cb(),
@@ -4342,6 +4495,14 @@ void llama_set_embeddings_nextn(llama_context * ctx, bool value, bool masked) {
ctx->set_embeddings_nextn(value, masked);
}
+void llama_set_embeddings_layer_inp_dev(llama_context * ctx, bool value) {
+ ctx->set_embeddings_layer_inp_dev(value);
+}
+
+bool llama_supports_embd_row(const llama_context * ctx) {
+ return ctx->supports_embd_row();
+}
+
void llama_set_embeddings_layer_inp(llama_context * ctx, uint32_t lid, bool value) {
ctx->set_embeddings_layer_inp(lid, value);
}
diff --git a/src/llama-context.h b/src/llama-context.h
index 6938fb103..39c5619df 100644
--- a/src/llama-context.h
+++ b/src/llama-context.h
@@ -93,6 +93,8 @@ struct llama_context {
float * get_embeddings_layer_inp(uint32_t lid);
+ ggml_tensor * get_layer_inp_dev() const { return layer_inp_dev; }
+
llama_token * get_sampled_tokens() const;
llama_token get_sampled_token_ith(int32_t idx);
@@ -120,6 +122,10 @@ struct llama_context {
void set_embeddings (bool value);
void set_embeddings_nextn(bool value, bool masked);
void set_embeddings_layer_inp(uint32_t lid, bool enable);
+ void set_embeddings_layer_inp_dev(bool value);
+
+ // whether decode accepts entries whose embd is a row of ctx_other's layer input device buffer
+ bool supports_embd_row() const;
void set_nextn_layer_offset(int32_t offset);
void set_causal_attn(bool value);
void set_warmup(bool value);
@@ -243,6 +249,11 @@ private:
// from backend into host-side embd_layer_inp buffers
bool extract_layer_inputs(const llm_graph_result * res, size_t token_offset, size_t n_tokens);
+ // (re)allocate the device buffer for cparams.embeddings_layer_inp_dev
+ void layer_inp_dev_alloc();
+
+ ggml_backend_t backend_from_dev(ggml_backend_dev_t dev) const;
+
//
// graph
//
@@ -420,6 +431,20 @@ private:
// host buffer for the model output (logits and embeddings)
ggml_backend_buffer_ptr buf_output;
+ // device buffer with enabled layer inputs, [n_layer_inp*n_embd, n_batch], row = batch index
+ ggml_context_ptr ctx_layer_inp_dev;
+ ggml_backend_buffer_ptr buf_layer_inp_dev;
+ ggml_tensor * layer_inp_dev = nullptr;
+ ggml_backend_dev_t layer_inp_dev_dev = nullptr;
+ ggml_backend_t layer_inp_backend = nullptr;
+
+ // ordering with the reader context (no host sync):
+ // written - recorded by this context after the batch, waited on by the reader
+ // read - recorded by the reader after it consumed the rows, waited on by this context
+ ggml_backend_event_ptr ev_layer_inp_written;
+ ggml_backend_event_ptr ev_layer_inp_read;
+ bool ev_layer_inp_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 c02e8bcac..9141cba3e 100644
--- a/src/llama-cparams.h
+++ b/src/llama-cparams.h
@@ -62,6 +62,7 @@ struct llama_cparams {
size_t moe_cache_size;
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
enum llama_context_type ctx_type;
enum llama_rope_scaling_type rope_scaling_type;
diff --git a/src/llama-ext.h b/src/llama-ext.h
index fcd91e912..23a781404 100644
--- a/src/llama-ext.h
+++ b/src/llama-ext.h
@@ -164,6 +164,19 @@ LLAMA_API void llama_set_embeddings_layer_inp(struct llama_context * ctx, uint32
// LLAMA_API float * llama_get_embeddings(struct llama_context * ctx);
LLAMA_API float * llama_get_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid);
+// Write the enabled layer inputs to a device buffer [n_layer_inp*n_embd, n_batch] instead of host memory.
+// Row r holds the token at batch index r, layers in ascending layer id order.
+LLAMA_API void llama_set_embeddings_layer_inp_dev(struct llama_context * ctx, bool value);
+
+// Whether decode accepts entries added with llama_batch_ext_add_embd_row()
+// (needs a DFlash/DSpark decoder context whose ctx_other has a layer input device buffer).
+LLAMA_API bool llama_supports_embd_row(const struct llama_context * ctx);
+
+// Add an entry whose embedding is row `row` of the layer input device buffer of ctx_other.
+// All entries of the batch must be of this kind. Positions are set with llama_batch_ext_set_pos(), n_pos_per_embd of them.
+// Returns the batch index (>= 0), or the llama_batch_ext_add() error codes.
+LLAMA_API int32_t llama_batch_ext_add_embd_row(struct llama_batch_ext * batch, llama_seq_id seq_id, int32_t row);
+
LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx);
//
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 725bdb048..fe8bdebe5 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -138,6 +138,12 @@ void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {
ggml_backend_tensor_set(tokens, ubatch->token, 0, n_tokens*ggml_element_size(tokens));
}
+ if (ubatch->embd_row && embd_rows) {
+ const int64_t n_tokens = ubatch->n_tokens;
+
+ ggml_backend_tensor_set(embd_rows, ubatch->embd_row, 0, n_tokens*ggml_element_size(embd_rows));
+ }
+
if (ubatch->embd && embd && !ubatch->is_mixed()) {
GGML_ASSERT(n_embd == embd->ne[0]);
@@ -197,6 +203,7 @@ bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) {
res &= (!mixed_tokens) || mixed_tokens->ne[0] == llm_graph_n_tok_rows(params.ubatch);
res &= (!mixed_embd) || mixed_embd->ne[1] == params.ubatch.n_tokens;
res &= (!scale_rows) || scale_rows->ne[1] == params.ubatch.n_tokens;
+ res &= (!params.ubatch.embd_row) || (embd_rows && embd_rows->ne[0] == params.ubatch.n_tokens);
return res;
}
@@ -232,6 +239,22 @@ bool llm_graph_input_embd_h::can_reuse(const llm_graph_params & params) {
return res;
}
+void llm_graph_input_batch_idxs::set_input(const llama_ubatch * ubatch) {
+ const int64_t n_tokens = ubatch->n_tokens;
+
+ GGML_ASSERT(ubatch->data && ubatch->data->batch_idxs.size() == (size_t) n_tokens);
+ GGML_ASSERT(ggml_backend_buffer_is_host(idxs->buffer));
+
+ int64_t * data = (int64_t *) idxs->data;
+ for (int64_t i = 0; i < n_tokens; ++i) {
+ data[i] = ubatch->data->batch_idxs[i];
+ }
+}
+
+bool llm_graph_input_batch_idxs::can_reuse(const llm_graph_params & params) {
+ return idxs->ne[0] == params.ubatch.n_tokens;
+}
+
void llm_graph_input_pos::set_input(const llama_ubatch * ubatch) {
if (ubatch->pos && pos) {
const int64_t n_tokens = ubatch->n_tokens;
@@ -1436,7 +1459,7 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) {
}
{
const auto & embeddings_layer_inp = params.cparams.embeddings_layer_inp;
- for (size_t il = 0; il < embeddings_layer_inp.size(); ++il) {
+ for (size_t il = 0; il < embeddings_layer_inp.size() && !params.cparams.embeddings_layer_inp_dev; ++il) {
if (embeddings_layer_inp[il]) {
GGML_ASSERT(t_layer_inp[il] != nullptr && "layer input tensor is null");
ggml_set_output(t_layer_inp[il]);
@@ -1563,6 +1586,8 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
moe_cache (params.moe_cache),
prec_policy (params.prec_policy),
hdmd (params.hdmd),
+ layer_inp_dev (params.layer_inp_dev),
+ layer_inp_src (params.layer_inp_src),
samplers (params.samplers),
cb_func (params.cb),
res (params.res),
@@ -4223,6 +4248,36 @@ void llm_graph_context::build_logits() const {
}
}
+void llm_graph_context::build_layer_inp_dev() const {
+ if (!cparams.embeddings_layer_inp_dev || layer_inp_dev == nullptr) {
+ return;
+ }
+
+ 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);
+
+ // row layout: [layer_0 | layer_1 | ...] in ascending layer id order
+ int64_t col = 0;
+ for (size_t il = 0; il < cparams.embeddings_layer_inp.size(); ++il) {
+ if (!cparams.embeddings_layer_inp[il]) {
+ continue;
+ }
+
+ ggml_tensor * src = res->t_layer_inp[il];
+ GGML_ASSERT(src != nullptr && "layer input tensor is null");
+ GGML_ASSERT(src->ne[0] == n_embd);
+
+ ggml_tensor * dst = ggml_view_2d(ctx0, layer_inp_dev, n_embd, layer_inp_dev->ne[1], layer_inp_dev->nb[1], col*n_embd*ggml_element_size(layer_inp_dev));
+ ggml_build_forward_expand(gf, ggml_set_rows(ctx0, dst, ggml_reshape_2d(ctx0, src, n_embd, n_tokens), inp->idxs));
+ col++;
+ }
+ GGML_ASSERT(col*n_embd == layer_inp_dev->ne[0]);
+
+ 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 b528ee11d..37b47acbb 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -147,6 +147,7 @@ public:
ggml_tensor * mixed_slots = nullptr; // I64 [n_tok_rows], mixed path: batch index of the token rows
ggml_tensor * mixed_embd = nullptr; // F32 [n_embd, n_batch], mixed path: embd rows, token rows are overwritten
ggml_tensor * scale_rows = nullptr; // F32 [1, n_batch], per-row scale: scale_tok for token rows, 1 for embd rows
+ ggml_tensor * embd_rows = nullptr; // I32 [n_batch], embd row path: rows of ctx_other's layer input device buffer
float scale_tok = 1.0f;
@@ -171,6 +172,19 @@ public:
const int64_t n_embd_state = 0;
};
+// original batch index of each ubatch token, used as set_rows destination rows
+class llm_graph_input_batch_idxs : public llm_graph_input_i {
+public:
+ llm_graph_input_batch_idxs() = default;
+ virtual ~llm_graph_input_batch_idxs() = default;
+
+ void set_input(const llama_ubatch * ubatch) override;
+
+ bool can_reuse(const llm_graph_params & params) override;
+
+ ggml_tensor * idxs = nullptr; // I64 [n_batch]
+};
+
class llm_graph_input_pos : public llm_graph_input_i {
public:
llm_graph_input_pos(uint32_t n_pos_per_embd, llama_rope_section_order rope_section_order = LLAMA_ROPE_SECTION_ORDER_TYXZ)
@@ -794,6 +808,9 @@ struct llm_graph_params {
const llama_hadamard * hdmd = nullptr;
+ 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)
+
std::map<llama_seq_id, llama_sampler *> samplers;
static bool samplers_equal(
@@ -833,7 +850,8 @@ struct llm_graph_params {
(!ubatch.embd && !other.ubatch.embd) ||
(ubatch.token && other.ubatch.token && ubatch.embd && other.ubatch.embd)
) &&
- (!ubatch.embd_state == !other.ubatch.embd_state);
+ (!ubatch.embd_state == !other.ubatch.embd_state) &&
+ (!ubatch.embd_row == !other.ubatch.embd_row);
// when we split the batch using "equal_seqs" we have to verify that the participating sequences are the same
// the reason is because the set of attention streams would be different for different sequences
@@ -885,6 +903,8 @@ struct llm_graph_params {
cparams.embeddings_nextn == other.cparams.embeddings_nextn &&
cparams.embeddings_nextn_masked == other.cparams.embeddings_nextn_masked &&
cparams.causal_attn == other.cparams.causal_attn &&
+ layer_inp_dev == other.layer_inp_dev &&
+ layer_inp_src == other.layer_inp_src &&
arch == other.arch &&
gtype == other.gtype &&
cvec == other.cvec &&
@@ -1052,6 +1072,9 @@ struct llm_graph_context {
const llama_hadamard * hdmd;
+ ggml_tensor * layer_inp_dev;
+ ggml_tensor * layer_inp_src;
+
std::map<llama_seq_id, llama_sampler *> samplers;
const llm_graph_cb & cb_func;
@@ -1450,6 +1473,9 @@ struct llm_graph_context {
void build_sampling() const;
+ // write enabled layer inputs into layer_inp_dev at their batch rows
+ void build_layer_inp_dev() const;
+
//
// dense (out)
//
diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp
index 10da95ae8..d5506f7a7 100644
--- a/src/llama-kv-cache-dsv4.cpp
+++ b/src/llama-kv-cache-dsv4.cpp
@@ -160,6 +160,7 @@ static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) {
/*.token =*/ data->token.empty() ? nullptr : data->token.data(),
/*.embd =*/ nullptr,
/*.embd_state =*/ nullptr,
+ /*.embd_row =*/ nullptr,
/*.pos =*/ data->pos.data(),
/*.n_seq_id =*/ data->n_seq_id.data(),
/*.seq_id =*/ data->seq_id.data(),
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 52893876d..7529f55c2 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -3400,6 +3400,8 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
// TODO: move reranking logic here and generalize
llm->build_dense_out(dense_2_out_layers, dense_2_out_layers_b, dense_3_out_layers);
+ llm->build_layer_inp_dev();
+
llm->res->set_outputs(params);
return llm->res->get_gf();
diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
index 209498421..fdd67e017 100644
--- a/src/models/dflash.cpp
+++ b/src/models/dflash.cpp
@@ -434,6 +434,23 @@ llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_grap
ggml_build_forward_expand(gf, cur);
}
+// target features of an injection batch: host embd rows, or rows of the target's layer input device buffer
+static ggml_tensor * build_inp_target_features(llm_graph_context & g, llm_graph_input_embd & inp, int64_t n_embd_inp) {
+ if (g.ubatch.embd_row) {
+ GGML_ASSERT(g.layer_inp_src != nullptr && g.layer_inp_src->ne[0] == n_embd_inp);
+
+ inp.embd_rows = ggml_new_tensor_1d(g.ctx0, GGML_TYPE_I32, g.n_tokens);
+ ggml_set_input(inp.embd_rows);
+
+ return ggml_get_rows(g.ctx0, g.layer_inp_src, inp.embd_rows);
+ }
+
+ inp.embd = ggml_new_tensor_2d(g.ctx0, GGML_TYPE_F32, n_embd_inp, g.n_tokens);
+ ggml_set_input(inp.embd);
+
+ return inp.embd;
+}
+
// DSpark (DFlash + Markov & Confidence head): Markov bias on the draft logits, chained per block position
static void build_dspark_markov_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens, ggml_tensor * hidden = nullptr) {
ggml_context * ctx0 = g.ctx0;
@@ -939,13 +956,10 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
// KV cache injection
ASSERT_EMBD_OR_TOKEN(ubatch);
- if (ubatch.embd) {
+ if (ubatch.embd || ubatch.embd_row) {
auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
- inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
- ggml_set_input(inp->embd);
-
- ggml_tensor * inp_target = inp->embd;
+ ggml_tensor * inp_target = build_inp_target_features(*this, *inp, n_embd_inp);
cb(inp_target, "inp_target_features", -1);
res->add_input(std::move(inp));
@@ -1219,13 +1233,10 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
// KV cache injection: fused target features from the encoder
ASSERT_EMBD_OR_TOKEN(ubatch);
- if (ubatch.embd) {
+ if (ubatch.embd || ubatch.embd_row) {
auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
- inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
- ggml_set_input(inp->embd);
-
- ggml_tensor * inp_target = inp->embd;
+ ggml_tensor * inp_target = build_inp_target_features(*this, *inp, n_embd_inp);
cb(inp_target, "inp_target_features", -1);
res->add_input(std::move(inp));
@@ -1418,13 +1429,10 @@ llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const ll
};
ASSERT_EMBD_OR_TOKEN(ubatch);
- if (ubatch.embd) {
+ if (ubatch.embd || ubatch.embd_row) {
auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
- inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
- ggml_set_input(inp->embd);
-
- ggml_tensor * inp_target = inp->embd;
+ ggml_tensor * inp_target = build_inp_target_features(*this, *inp, n_embd_inp);
cb(inp_target, "inp_target_features", -1);
res->add_input(std::move(inp));
@@ -1668,13 +1676,10 @@ llama_model_dflash::graph_mistral_large4::graph_mistral_large4(const llama_model
};
ASSERT_EMBD_OR_TOKEN(ubatch);
- if (ubatch.embd) {
+ if (ubatch.embd || ubatch.embd_row) {
auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
- inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
- ggml_set_input(inp->embd);
-
- ggml_tensor * inp_target = inp->embd;
+ ggml_tensor * inp_target = build_inp_target_features(*this, *inp, n_embd_inp);
cb(inp_target, "inp_target_features", -1);
res->add_input(std::move(inp));