Commit 32bdc3a89 for llama.cpp
commit 32bdc3a89e6528d3e419024c833c5dabf32c7541
Author: Ruixiang Wang <wangruixiang07@outlook.com>
Date: Sat Nov 14 09:49:27 2026 +0100
spec + model: refactor DFlash based speculative decoding methods (#32016)
* dflash : move the DSpark markov head and backbone graphs to dspark.cpp
Move-only: build_dspark_markov_head, graph_dsv4, graph_dsv41 and
graph_mistral_large4 move verbatim. build_inp_target_features and
build_dspark_markov_head become static members of llama_model_dflash so
both files can use them.
* dflash : detect the head and the backbone once, split the loaders per file
load_arch_hparams sets llama_model_dflash::variant (DFlash, DFlash2,
DSpark, LiliCorr) and ::backbone (base, DSV4, DSV4.1, Mistral Large 4).
The head, backbone and base layer loaders become member helpers in
dflash.cpp, dflash2.cpp, lilicorr.cpp and dspark.cpp. build_arch_graph
dispatches on the backbone.
The DFlash2 and LiliCorr loaders create the DFlash2 conv tensors through
a helper in dflash2.cpp, right after their head tensors instead of in
the base layer loop; the other tensors keep their creation order. The
DFlash2 loader drops its copy of the conv group size check, which master
never reached.
* dflash : share the injection input, token embeddings and lm head lookups
build_inject_inp (target features + fc + encoder norm), get_tok_embd and
get_lm_head replace the copies in the base, DSV4, DSV4.1 and Mistral
Large 4 graphs. build_inp_target_features is file-static again.
* dflash : split the DFlash decoder into a graph per head
graph_base creates the shared inputs and provides build_inject,
build_inp_tokens, build_layers and build_logits. graph_dflash
(dflash.cpp), graph_dflash2 (dflash2.cpp), graph_dspark (dspark.cpp) and
graph_lilicorr (lilicorr.cpp) assemble them and add their head; the
selector, scorer and markov head builders move to the file of their
graph, and the DFlash2 local conv run by build_layers to dflash2.cpp.
build_arch_graph dispatches on the backbone, then on the head.
The op order and the tensor names are unchanged.
* speculative : split the DFlash driver into a class per head
common_speculative_impl_draft_dflash stays the DFlash implementation and
becomes the base of draft_dspark, draft_dflash2 and draft_lilicorr. Its
protected constructor only reads the draft metadata; init() then sets up
the contexts with the options of the draft (draft length, backend
sampling, device features, masked h_nextn). Each derived class overrides
draft() with the path of its head from the former draft(). DFlash2 also
overrides begin() to reseed its generator per request, and LiliCorr
overrides process() to store the context features of each injected
chunk, which its noise block carries as its state.
The factory keeps the current routing: draft-dspark picks DSpark,
draft-dflash picks DFlash2 or LiliCorr from the draft metadata and
DFlash otherwise. draft-dspark with a DFlash2 or LiliCorr draft now runs
the DSpark reader instead of a mix of both. sample_from_anchor is logged
by DSpark only.
* speculative : add the draft-dflash2 and draft-lilicorr types
draft-dflash2 and draft-lilicorr name the DFlash2 and LiliCorr drafts,
which draft-dflash keeps running too: it picks them from the draft
metadata. draft-dflash2 and draft-lilicorr must match the draft model,
draft-dspark rejects a DFlash2 or LiliCorr draft, a model that is not a
DFlash draft is rejected, and a run uses at most one DFlash family type;
these errors leave speculative decoding uninitialized.
The type detected from the draft GGUF is now the concrete one, and
draft-dflash2 and draft-lilicorr download the dflash sidecar.
* docs : describe the DFlash2 and LiliCorr drafts
Add the DFlash2 and LiliCorr sections and the new types, and explain
that draft-dflash also runs DFlash2 and LiliCorr drafts.
* speculative : drop the unused DFlash is_mrope
Nothing reads it since the migration to common_batch (#29385), which
picks the number of M-RoPE position rows from the rope type of the draft.
* doc: update spec dec
* address review comments
diff --git a/common/arg.cpp b/common/arg.cpp
index 6bf9355fb..00f9659cb 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -349,6 +349,14 @@ common_models_handler common_models_handler_init(const common_params & params, l
params.speculative.types.end(),
COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != params.speculative.types.end();
+ const bool spec_type_draft_dflash2 = std::find(params.speculative.types.begin(),
+ params.speculative.types.end(),
+ COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2) != params.speculative.types.end();
+
+ const bool spec_type_draft_lilicorr = std::find(params.speculative.types.begin(),
+ params.speculative.types.end(),
+ COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR) != params.speculative.types.end();
+
const bool spec_type_draft_eagle3 = std::find(params.speculative.types.begin(),
params.speculative.types.end(),
COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3) != params.speculative.types.end();
@@ -373,7 +381,7 @@ common_models_handler common_models_handler_init(const common_params & params, l
opts.offline = params.offline;
opts.download_mtp = spec_type_draft_mtp;
opts.download_eagle3 = spec_type_draft_eagle3;
- opts.download_dflash = spec_type_draft_dflash;
+ opts.download_dflash = spec_type_draft_dflash || spec_type_draft_dflash2 || spec_type_draft_lilicorr;
opts.download_dspark = spec_type_draft_dspark;
opts.download_mmproj = use_mmproj && !params.no_mmproj
&& params.mmproj.path.empty() && params.mmproj.url.empty();
diff --git a/common/common.h b/common/common.h
index 6043ab342..4cfc1b8f3 100644
--- a/common/common.h
+++ b/common/common.h
@@ -174,18 +174,20 @@ enum common_params_sampling_config : uint64_t {
};
enum common_speculative_type {
- COMMON_SPECULATIVE_TYPE_NONE, // no speculative decoding
- COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, // standalone draft model speculative decoding
- COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, // Eagle3 speculative decoding
- COMMON_SPECULATIVE_TYPE_DRAFT_MTP, // Multi-token prediction
- COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, // DFlash speculative decoding
- COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, // DSpark speculative decoding (DFlash + Markov head)
- COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding based on n-grams
- COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only
- COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values
+ COMMON_SPECULATIVE_TYPE_NONE, // no speculative decoding
+ COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, // standalone draft model speculative decoding
+ COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, // Eagle3 speculative decoding
+ COMMON_SPECULATIVE_TYPE_DRAFT_MTP, // Multi-token prediction
+ COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, // DFlash speculative decoding
+ COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, // DSpark speculative decoding (DFlash + Markov head)
+ COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2, // DFlash2 speculative decoding (DFlash + selector)
+ COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR, // LiliCorr speculative decoding (DFlash + scorer)
+ COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding based on n-grams
+ COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only
+ COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values
COMMON_SPECULATIVE_TYPE_NGRAM_MOD,
- COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, // self-speculative decoding with 3-level n-gram cache
- COMMON_SPECULATIVE_TYPE_COUNT // number of types, unknown type
+ COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, // self-speculative decoding with 3-level n-gram cache
+ COMMON_SPECULATIVE_TYPE_COUNT // number of types, unknown type
};
// Grammar type enumeration
@@ -414,7 +416,12 @@ struct common_params_speculative {
uint32_t need_n_rs_seq() const {
bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) {
- return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
+ return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP ||
+ t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 ||
+ t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH ||
+ t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK ||
+ t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2 ||
+ t == COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR;
});
return needs_rs_seq ? draft.n_max : 0u;
diff --git a/common/speculative.cpp b/common/speculative.cpp
index 63cb5e611..b40eb5535 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -70,17 +70,19 @@ static void spec_retune(
}
const std::map<std::string, common_speculative_type> common_speculative_type_from_name_map = {
- {"none", COMMON_SPECULATIVE_TYPE_NONE},
- {"draft-simple", COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE},
- {"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3},
- {"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP},
- {"draft-dflash", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH},
- {"draft-dspark", COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK},
- {"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE},
- {"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},
- {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},
- {"ngram-mod", COMMON_SPECULATIVE_TYPE_NGRAM_MOD},
- {"ngram-cache", COMMON_SPECULATIVE_TYPE_NGRAM_CACHE}
+ {"none", COMMON_SPECULATIVE_TYPE_NONE},
+ {"draft-simple", COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE},
+ {"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3},
+ {"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP},
+ {"draft-dflash", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH},
+ {"draft-dspark", COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK},
+ {"draft-dflash2", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2},
+ {"draft-lilicorr", COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR},
+ {"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE},
+ {"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},
+ {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},
+ {"ngram-mod", COMMON_SPECULATIVE_TYPE_NGRAM_MOD},
+ {"ngram-cache", COMMON_SPECULATIVE_TYPE_NGRAM_CACHE}
};
static std::string common_speculative_get_devices_str(const std::vector<ggml_backend_dev_t> & devices) {
@@ -999,26 +1001,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
int32_t block_size = 0;
llama_token mask_token_id = 0;
- bool is_lilicorr = false;
- int32_t lilicorr_top_k = 0;
- size_t lilicorr_rows = 0;
- std::vector<float> lilicorr_embd;
- std::vector<llama_pos> lilicorr_pos;
-
- bool is_dflash2 = false;
- bool is_mrope = false;
- int32_t selector_top_k = 0;
-
- // DFlash2 probabilistic drafting: one generator per sequence and the seed it was last seeded with (-1: reseed)
- std::vector<std::mt19937> dflash2_rng;
- std::vector<int64_t> dflash2_seed;
-
- // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
- const bool is_dspark;
-
- // dspark speculators
- bool sample_from_anchor = true;
-
// block-internal attention
bool causal_attn = false;
@@ -1028,11 +1010,31 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
// 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_impl_draft_dflash(params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
+ {
+ // DFlash input is [id_last, <mask> * (block_size-1)]: in-place denoising yields at most block_size-1 draft tokens
+ init({
+ /* .n_draft_max = */ block_size - 1,
+ /* .backend_sampling = */ true,
+ /* .dev_features = */ true,
+ /* .nextn_masked = */ true,
+ });
+ }
+
+protected:
+ // what init() sets up, per draft type
+ struct dflash_options {
+ int32_t n_draft_max; // longest draft a block yields, n_max and n_min are clamped to it
+ bool backend_sampling; // the block logits can be sampled on the backend
+ bool dev_features; // the target features can come from the target's device buffer
+ bool nextn_masked; // h_nextn is only read at the output rows
+ };
+
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_type type)
: common_speculative_impl(type, n_seq, params.draft.n_max)
, params(params.draft)
- , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)
{
auto * ctx_tgt = this->params.ctx_tgt;
auto * ctx_dft = this->params.ctx_dft;
@@ -1056,48 +1058,26 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
if (llama_model_meta_val_str(model_dft, "dflash.block_size", buf, sizeof(buf)) >= 0) {
block_size = std::atoi(buf);
}
- if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) {
- sample_from_anchor = std::strcmp(buf, "true") == 0;
- }
if (llama_model_meta_val_str(model_dft, "dflash.attention.causal", buf, sizeof(buf)) >= 0) {
causal_attn = std::strcmp(buf, "true") == 0;
}
}
- selector_top_k = llama_model_dflash_selector_top_k(model_dft);
- is_dflash2 = selector_top_k > 0;
- dflash2_rng.resize(n_seq);
- dflash2_seed.assign(n_seq, -1);
- char lilicorr_meta[32] = {};
- lilicorr_top_k = llama_model_meta_val_str(model_dft, "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0
- ? std::atoi(lilicorr_meta) : 0;
- is_lilicorr = lilicorr_top_k > 0;
- if (is_lilicorr) {
- lilicorr_rows = (size_t) block_size + 1;
- const size_t lilicorr_slots = (size_t) n_seq * lilicorr_rows;
- lilicorr_embd.resize(lilicorr_slots * (size_t) n_embd_dec);
- lilicorr_pos.assign(lilicorr_slots, -1);
- }
mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
- if (is_dspark && this->params.p_min > 0.0f) {
- char buf[16] = {};
- const bool has_conf =
- llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 ||
- std::strcmp(buf, "true") == 0;
- if (!has_conf) {
- throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0");
- }
- }
-
LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str());
LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);
- LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__,
- block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false");
+ LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u\n", __func__,
+ block_size, mask_token_id, target_layer_ids_n);
+ }
+
+ void init(const dflash_options & opt) {
+ auto * ctx_tgt = this->params.ctx_tgt;
+ auto * ctx_dft = this->params.ctx_dft;
- // DFlash input is [id_last, <mask> * (block_size-1)]: in-place denoising yields at most
- // block_size-1 draft tokens, anchor-first DSpark yields a full block_size draft tokens
- const int32_t n_draft_max = is_dspark && sample_from_anchor ? block_size : block_size - 1;
+ const llama_model * model_dft = llama_get_model(ctx_dft);
+
+ const int32_t n_draft_max = opt.n_draft_max;
if (this->params.n_max > n_draft_max || this->params.n_min > n_draft_max) {
LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained block size %d -- clamping to %d\n",
__func__, this->params.n_max, this->params.n_min, block_size, n_draft_max);
@@ -1109,9 +1089,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
batch = common_batch(ctx_dft);
batch_inject = common_batch(ctx_dft);
- // embd batches on an M-RoPE draft carry 4 position rows per token
- is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE;
-
smpls.resize(n_seq);
for (auto & s : smpls) {
common_params_sampling sparams;
@@ -1123,7 +1100,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
// offload draft sampling to the backend
backend_chains.assign(n_seq, nullptr);
- if (this->params.backend_sampling && !is_dflash2 && !is_lilicorr) {
+ if (this->params.backend_sampling && opt.backend_sampling) {
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params());
llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));
@@ -1148,7 +1125,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
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);
+ use_dev_features = opt.dev_features && 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);
@@ -1159,11 +1136,15 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
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_embeddings_nextn(ctx_dft, true, opt.nextn_masked);
llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise
}
+ // only LiliCorr uses this, to keep the context features for its scorer; the other heads only need the KV cache
+ // called after the host copy injection of a chunk, while h_nextn still holds the context features of batch_inject
+ virtual void store_context_features(llama_seq_id /*seq_id*/, int32_t /*n_chunk*/) {}
+
+public:
~common_speculative_impl_draft_dflash() override {
auto * ctx_dft = this->params.ctx_dft;
for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) {
@@ -1183,9 +1164,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
return;
}
- // reseed per request rather than per round, or two identical requests differ
- dflash2_seed[seq_id] = -1;
-
const int32_t N = (int32_t) prompt.size();
if (N <= 0) {
return;
@@ -1297,24 +1275,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
__func__, rc, (int) n_chunk, (int) offset);
return false;
}
- if (is_lilicorr) {
- // Ordinary DFlash can write to the KV cache, but LiliCorr needs to store embeddings across sequences.
- const float * projected = llama_get_embeddings_nextn(ctx_dft);
- if (projected) {
- const size_t seq_offset = (size_t) seq_id * lilicorr_rows;
- const llama_pos pos_min = std::max(0, batch_inject.tokens.back().pos[0] - block_size);
- for (int32_t i = 0; i < n_chunk; ++i) {
- const llama_pos pos = batch_inject.tokens[i].pos[0];
- if (pos < pos_min) {
- continue;
- }
- const size_t slot = seq_offset + (size_t) pos % lilicorr_rows;
- std::memcpy(lilicorr_embd.data() + slot * (size_t) n_embd_dec,
- projected + (size_t) i * n_embd_dec, (size_t) n_embd_dec * sizeof(float));
- lilicorr_pos[slot] = pos;
- }
- }
- }
+ store_context_features(seq_id, n_chunk);
}
}
@@ -1326,7 +1287,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
batch.clear();
- // build one batch holding every drafting sequence's noise block into a single decode)
+ // build one batch holding every drafting sequence's noise block into a single decode
// record where each block starts and its size
std::vector<int32_t> i_block_beg(n_seq, -1);
std::vector<int32_t> n_block (n_seq, 0);
@@ -1342,17 +1303,136 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
const int32_t n = (int32_t) dp.pos0;
const int32_t n_draft = params.n_max;
- const int32_t n_block_tokens = is_lilicorr ? block_size : n_draft + (is_dspark && sample_from_anchor ? 0 : 1);
+ const int32_t n_block_tokens = n_draft + 1;
i_block_beg[seq_id] = batch.size();
n_block [seq_id] = n_block_tokens;
for (int32_t i = 0; i < n_block_tokens; ++i) {
- const int32_t idx = batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, !is_dflash2 && !is_lilicorr);
- if (is_lilicorr) {
- // For LiliCorr drafting batch, set state embeddings for use in the associated transformer head.
- const llama_pos pos = dp.pos0 - 1;
- const size_t slot = (size_t) seq_id * lilicorr_rows + (size_t) pos % lilicorr_rows;
- batch.set_embd_state(idx, { lilicorr_embd.data() + slot * (size_t) n_embd_dec, 1, (size_t) n_embd_dec });
+ batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, true);
+ }
+ }
+
+ if (batch.size() == 0) {
+ return;
+ }
+
+ // decode all sequence's noise block in a single batch
+ int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+ if (ret != 0) {
+ LOG_WRN("%s: llama_process returned %d\n", __func__, ret);
+ return;
+ }
+
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ if (i_block_beg[seq_id] < 0) {
+ continue;
+ }
+ auto & dp = dparams[seq_id];
+
+ const int32_t beg = i_block_beg[seq_id];
+ const int32_t n_block_tokens = n_block[seq_id];
+
+ auto * smpl = smpls[seq_id].get();
+
+ auto & result = *dp.result;
+
+ // greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1
+ for (int32_t i = 1; i < n_block_tokens; ++i) {
+ common_sampler_sample(smpl, ctx_dft, beg + i, true);
+
+ const auto * cur_p = common_sampler_get_candidates(smpl, true);
+
+ for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
+ LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
+ seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p,
+ common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
+ }
+
+ const llama_token id = cur_p->data[0].id;
+
+ if (cur_p->data[0].p < params.p_min) {
+ break;
}
+
+ common_sampler_accept(smpl, id, true);
+
+ result.push_back(id);
+ }
+
+ if (result.size() < (size_t) params.n_min) {
+ result.clear();
+ }
+ }
+ }
+
+ void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
+ // noop
+ }
+};
+
+// DSpark: a semi-autoregressive Markov head (+ Confidence head) on the block, anchor-first block layout
+struct common_speculative_impl_draft_dspark final : public common_speculative_impl_draft_dflash {
+ // bonus anchor (SpecForge exports) when false: slot 0 is a bonus token, not a prediction slot
+ bool sample_from_anchor = true;
+
+ common_speculative_impl_draft_dspark(const common_params_speculative & params, uint32_t n_seq)
+ : common_speculative_impl_draft_dflash(params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)
+ {
+ const llama_model * model_dft = llama_get_model(this->params.ctx_dft);
+
+ {
+ char buf[32] = {};
+ if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) {
+ sample_from_anchor = std::strcmp(buf, "true") == 0;
+ }
+ }
+
+ if (this->params.p_min > 0.0f) {
+ char buf[16] = {};
+ const bool has_conf =
+ llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 ||
+ std::strcmp(buf, "true") == 0;
+ if (!has_conf) {
+ throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0");
+ }
+ }
+
+ LOG_INF("%s: - sample_from_anchor=%s\n", __func__, sample_from_anchor ? "true" : "false");
+
+ // anchor-first DSpark yields a full block_size draft tokens
+ init({
+ /* .n_draft_max = */ sample_from_anchor ? block_size : block_size - 1,
+ /* .backend_sampling = */ true,
+ /* .dev_features = */ true,
+ /* .nextn_masked = */ true,
+ });
+ }
+
+ void draft(common_speculative_draft_params_vec & dparams) override {
+ auto & ctx_dft = params.ctx_dft;
+
+ batch.clear();
+
+ // build one batch holding every drafting sequence's noise block into a single decode
+ // record where each block starts and its size
+ std::vector<int32_t> i_block_beg(n_seq, -1);
+ std::vector<int32_t> n_block (n_seq, 0);
+
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ auto & dp = dparams[seq_id];
+ if (!dp.drafting) {
+ continue;
+ }
+
+ common_sampler_reset(smpls[seq_id].get());
+
+ const int32_t n = (int32_t) dp.pos0;
+
+ const int32_t n_draft = params.n_max;
+ const int32_t n_block_tokens = n_draft + (sample_from_anchor ? 0 : 1);
+ i_block_beg[seq_id] = batch.size();
+ n_block [seq_id] = n_block_tokens;
+ for (int32_t i = 0; i < n_block_tokens; ++i) {
+ batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, true);
}
}
@@ -1380,136 +1460,327 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
auto & result = *dp.result;
- if (is_dflash2 || is_lilicorr) {
- // Both DFlash and LiliCorr greedily walk the candidate lattice.
- const float * lattice = llama_get_embeddings_nextn(ctx_dft);
- GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
+ // DSpark: read from the first draft slot, truncate below the confidence threshold
+ const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr;
+ // bonus-anchor drafts read the mask positions only, like DFlash
+ const int32_t i_draft_beg = sample_from_anchor ? 0 : 1;
+ for (int32_t i = i_draft_beg; i < n_block_tokens; ++i) {
+ const int32_t idx = beg + i;
+
+ if (conf && conf[(size_t) idx * n_embd_dec] < params.p_min) {
+ break;
+ }
- // DFlash2 can instead sample each step and hand q to the verifier for rejection sampling
- const bool sample = is_dflash2 && params.probabilistic && dp.result_q;
+ common_sampler_sample(smpl, ctx_dft, idx, true);
- // greedy keeps temperature 1, sampling uses the target's
- const float temp = sample ? dp.temp : 1.0f;
+ const auto * cur_p = common_sampler_get_candidates(smpl, true);
- if (sample && dflash2_seed[seq_id] != dp.seed) {
- dflash2_seed[seq_id] = dp.seed;
- // the default seed means random; mixed so it differs from the target's, as in spec_retune
- dflash2_rng[seq_id].seed(dp.seed == LLAMA_DEFAULT_SEED ? std::random_device{}() : dp.seed ^ 0x85ebca6bu);
+ for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
+ LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
+ seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
+ common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
}
- std::vector<float> probs;
- std::vector<int32_t> sorted_idx;
-
- int32_t predecessor = 0;
- for (int32_t i = 1; i < (is_dflash2 ? n_block_tokens : params.n_max + 1); ++i) {
- const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
- const int32_t top_k = is_dflash2 ? selector_top_k : lilicorr_top_k;
- const float * scores = row + top_k + (size_t) predecessor * top_k;
-
- predecessor = (int32_t) std::distance(scores,
- std::max_element(scores, scores + top_k));
- if (params.p_min > 0.0f || sample) {
- // softmax(scores / temp) at the argmax, i.e. 1 / sum(exp((s_k - s_max) / temp))
- probs.resize(top_k);
- float sum = 0.0f;
- for (int32_t k = 0; k < top_k; ++k) {
- probs[k] = std::exp((scores[k] - scores[predecessor]) / temp);
- sum += probs[k];
- }
- if (1.0f / sum < params.p_min) {
+ const llama_token id = cur_p->data[0].id;
+
+ common_sampler_accept(smpl, id, true);
+
+ result.push_back(id);
+ }
+
+ if (result.size() < (size_t) params.n_min) {
+ result.clear();
+ }
+ }
+ }
+};
+
+// DFlash2: walks the lattice of the selector over the top-k candidates of the block
+struct common_speculative_impl_draft_dflash2 final : public common_speculative_impl_draft_dflash {
+ int32_t selector_top_k = 0;
+
+ // probabilistic drafting: one generator per sequence and the seed it was last seeded with (-1: reseed)
+ std::vector<std::mt19937> dflash2_rng;
+ std::vector<int64_t> dflash2_seed;
+
+ common_speculative_impl_draft_dflash2(const common_params_speculative & params, uint32_t n_seq)
+ : common_speculative_impl_draft_dflash(params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2)
+ {
+ selector_top_k = llama_model_dflash_selector_top_k(llama_get_model(this->params.ctx_dft));
+ dflash2_rng.resize(n_seq);
+ dflash2_seed.assign(n_seq, -1);
+
+ // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits
+ init({
+ /* .n_draft_max = */ block_size - 1,
+ /* .backend_sampling = */ false,
+ /* .dev_features = */ true,
+ /* .nextn_masked = */ false,
+ });
+ }
+
+ void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {
+ if (seq_id >= 0 && seq_id < (llama_seq_id) n_seq) {
+ // reseed per request rather than per round, or two identical requests differ
+ dflash2_seed[seq_id] = -1;
+ }
+
+ common_speculative_impl_draft_dflash::begin(seq_id, prompt);
+ }
+
+ void draft(common_speculative_draft_params_vec & dparams) override {
+ auto & ctx_dft = params.ctx_dft;
+
+ batch.clear();
+
+ // build one batch holding every drafting sequence's noise block into a single decode
+ // record where each block starts and its size
+ std::vector<int32_t> i_block_beg(n_seq, -1);
+ std::vector<int32_t> n_block (n_seq, 0);
+
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ auto & dp = dparams[seq_id];
+ if (!dp.drafting) {
+ continue;
+ }
+
+ common_sampler_reset(smpls[seq_id].get());
+
+ const int32_t n = (int32_t) dp.pos0;
+
+ const int32_t n_draft = params.n_max;
+ const int32_t n_block_tokens = n_draft + 1;
+ i_block_beg[seq_id] = batch.size();
+ n_block [seq_id] = n_block_tokens;
+ for (int32_t i = 0; i < n_block_tokens; ++i) {
+ batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, false);
+ }
+ }
+
+ if (batch.size() == 0) {
+ return;
+ }
+
+ // decode all sequence's noise block in a single batch
+ int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+ if (ret != 0) {
+ LOG_WRN("%s: llama_process returned %d\n", __func__, ret);
+ return;
+ }
+
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ if (i_block_beg[seq_id] < 0) {
+ continue;
+ }
+ auto & dp = dparams[seq_id];
+
+ const int32_t beg = i_block_beg[seq_id];
+ const int32_t n_block_tokens = n_block[seq_id];
+
+ auto & result = *dp.result;
+
+ const float * lattice = llama_get_embeddings_nextn(ctx_dft);
+ GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
+
+ // DFlash2 can instead sample each step and hand q to the verifier for rejection sampling
+ const bool sample = params.probabilistic && dp.result_q;
+
+ // greedy keeps temperature 1, sampling uses the target's
+ const float temp = sample ? dp.temp : 1.0f;
+
+ if (sample && dflash2_seed[seq_id] != dp.seed) {
+ dflash2_seed[seq_id] = dp.seed;
+ // the default seed means random; mixed so it differs from the target's, as in spec_retune
+ dflash2_rng[seq_id].seed(dp.seed == LLAMA_DEFAULT_SEED ? std::random_device{}() : dp.seed ^ 0x85ebca6bu);
+ }
+
+ std::vector<float> probs;
+ std::vector<int32_t> sorted_idx;
+
+ int32_t predecessor = 0;
+ for (int32_t i = 1; i < n_block_tokens; ++i) {
+ const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
+ const int32_t top_k = selector_top_k;
+ const float * scores = row + top_k + (size_t) predecessor * top_k;
+
+ predecessor = (int32_t) std::distance(scores,
+ std::max_element(scores, scores + top_k));
+ if (params.p_min > 0.0f || sample) {
+ // softmax(scores / temp) at the argmax, i.e. 1 / sum(exp((s_k - s_max) / temp))
+ probs.resize(top_k);
+ float sum = 0.0f;
+ for (int32_t k = 0; k < top_k; ++k) {
+ probs[k] = std::exp((scores[k] - scores[predecessor]) / temp);
+ sum += probs[k];
+ }
+ if (1.0f / sum < params.p_min) {
+ break;
+ }
+ for (int32_t k = 0; k < top_k; ++k) {
+ probs[k] /= sum;
+ }
+ }
+ if (sample) {
+ // ggml_top_k leaves candidates unordered, so walk them by token id to keep a seed's draft stable
+ sorted_idx.resize(top_k);
+ for (int32_t k = 0; k < top_k; ++k) {
+ sorted_idx[k] = k;
+ }
+ std::sort(sorted_idx.begin(), sorted_idx.end(), [row](int32_t a, int32_t b) { return row[a] < row[b]; });
+
+ // a round-off overshoot keeps the argmax, which has q > 0
+ double u = std::uniform_real_distribution<double>(0.0, 1.0)(dflash2_rng[seq_id]);
+ for (const int32_t k : sorted_idx) {
+ u -= probs[k];
+ if (u < 0.0) {
+ predecessor = k;
break;
}
- for (int32_t k = 0; k < top_k; ++k) {
- probs[k] /= sum;
- }
}
- if (sample) {
- // ggml_top_k leaves candidates unordered, so walk them by token id to keep a seed's draft stable
- sorted_idx.resize(top_k);
- for (int32_t k = 0; k < top_k; ++k) {
- sorted_idx[k] = k;
- }
- std::sort(sorted_idx.begin(), sorted_idx.end(), [row](int32_t a, int32_t b) { return row[a] < row[b]; });
-
- // a round-off overshoot keeps the argmax, which has q > 0
- double u = std::uniform_real_distribution<double>(0.0, 1.0)(dflash2_rng[seq_id]);
- for (const int32_t k : sorted_idx) {
- u -= probs[k];
- if (u < 0.0) {
- predecessor = k;
- break;
- }
- }
- auto & q = dp.result_q->emplace_back();
- q.reserve(top_k);
- for (int32_t k = 0; k < top_k; ++k) {
- q.push_back({ (llama_token) row[k], scores[k], probs[k] });
- }
+ auto & q = dp.result_q->emplace_back();
+ q.reserve(top_k);
+ for (int32_t k = 0; k < top_k; ++k) {
+ q.push_back({ (llama_token) row[k], scores[k], probs[k] });
}
- result.push_back((llama_token) row[predecessor]);
}
+ result.push_back((llama_token) row[predecessor]);
+ }
+
+ if (result.size() < (size_t) params.n_min) {
+ result.clear();
+ if (sample) {
+ dp.result_q->clear();
+ }
+ }
+ }
+ }
+};
- if (result.size() < (size_t) params.n_min) {
- result.clear();
- if (sample) {
- dp.result_q->clear();
- }
+// LiliCorr: walks the lattice of the scorer over the top-k candidates of the block, which reads the
+// context features of the anchor; they come back from the injection and are kept per sequence
+struct common_speculative_impl_draft_lilicorr final : public common_speculative_impl_draft_dflash {
+ int32_t lilicorr_top_k = 0;
+
+ // context features of the last lilicorr_rows positions of each sequence, indexed by position
+ size_t lilicorr_rows = 0;
+ std::vector<float> lilicorr_embd;
+ std::vector<llama_pos> lilicorr_pos;
+
+ common_speculative_impl_draft_lilicorr(const common_params_speculative & params, uint32_t n_seq)
+ : common_speculative_impl_draft_dflash(params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR)
+ {
+ char lilicorr_meta[32] = {};
+ lilicorr_top_k = llama_model_meta_val_str(llama_get_model(this->params.ctx_dft), "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0
+ ? std::atoi(lilicorr_meta) : 0;
+ lilicorr_rows = (size_t) block_size + 1;
+ const size_t lilicorr_slots = (size_t) n_seq * lilicorr_rows;
+ lilicorr_embd.resize(lilicorr_slots * (size_t) n_embd_dec);
+ lilicorr_pos.assign(lilicorr_slots, -1);
+
+ // the context features are read back from h_nextn after a host copy injection
+ init({
+ /* .n_draft_max = */ block_size - 1,
+ /* .backend_sampling = */ false,
+ /* .dev_features = */ false,
+ /* .nextn_masked = */ false,
+ });
+ }
+
+protected:
+ void store_context_features(llama_seq_id seq_id, int32_t n_chunk) override {
+ // Ordinary DFlash can write to the KV cache, but LiliCorr needs to store embeddings across sequences.
+ const float * projected = llama_get_embeddings_nextn(params.ctx_dft);
+ if (projected) {
+ const size_t seq_offset = (size_t) seq_id * lilicorr_rows;
+ const llama_pos pos_min = std::max(0, batch_inject.tokens.back().pos[0] - block_size);
+ for (int32_t i = 0; i < n_chunk; ++i) {
+ const llama_pos pos = batch_inject.tokens[i].pos[0];
+ if (pos < pos_min) {
+ continue;
}
- continue;
+ const size_t slot = seq_offset + (size_t) pos % lilicorr_rows;
+ std::memcpy(lilicorr_embd.data() + slot * (size_t) n_embd_dec,
+ projected + (size_t) i * n_embd_dec, (size_t) n_embd_dec * sizeof(float));
+ lilicorr_pos[slot] = pos;
}
+ }
+ }
- if (is_dspark) {
- // DSpark: read from the first draft slot, truncate below the confidence threshold
- const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr;
- // bonus-anchor drafts read the mask positions only, like DFlash
- const int32_t i_draft_beg = sample_from_anchor ? 0 : 1;
- for (int32_t i = i_draft_beg; i < n_block_tokens; ++i) {
- const int32_t idx = beg + i;
+public:
+ void draft(common_speculative_draft_params_vec & dparams) override {
+ auto & ctx_dft = params.ctx_dft;
- if (conf && conf[(size_t) idx * n_embd_dec] < params.p_min) {
- break;
- }
+ batch.clear();
- common_sampler_sample(smpl, ctx_dft, idx, true);
+ // build one batch holding every drafting sequence's noise block into a single decode
+ // record where each block starts, all blocks have the same size
+ std::vector<int32_t> i_block_beg(n_seq, -1);
- const auto * cur_p = common_sampler_get_candidates(smpl, true);
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ auto & dp = dparams[seq_id];
+ if (!dp.drafting) {
+ continue;
+ }
- for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
- LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
- seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
- common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
- }
+ common_sampler_reset(smpls[seq_id].get());
- const llama_token id = cur_p->data[0].id;
+ const int32_t n = (int32_t) dp.pos0;
- common_sampler_accept(smpl, id, true);
+ // the scorer reads a full block
+ i_block_beg[seq_id] = batch.size();
+ for (int32_t i = 0; i < block_size; ++i) {
+ const int32_t idx = batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, false);
+ // For LiliCorr drafting batch, set state embeddings for use in the associated transformer head.
+ const llama_pos pos = dp.pos0 - 1;
+ const size_t slot = (size_t) seq_id * lilicorr_rows + (size_t) pos % lilicorr_rows;
+ batch.set_embd_state(idx, { lilicorr_embd.data() + slot * (size_t) n_embd_dec, 1, (size_t) n_embd_dec });
+ }
+ }
- result.push_back(id);
- }
- } else {
- // greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1
- for (int32_t i = 1; i < n_block_tokens; ++i) {
- common_sampler_sample(smpl, ctx_dft, beg + i, true);
+ if (batch.size() == 0) {
+ return;
+ }
- const auto * cur_p = common_sampler_get_candidates(smpl, true);
+ // decode all sequence's noise block in a single batch
+ int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+ if (ret != 0) {
+ LOG_WRN("%s: llama_process returned %d\n", __func__, ret);
+ return;
+ }
- for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
- LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
- seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p,
- common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
- }
+ for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+ if (i_block_beg[seq_id] < 0) {
+ continue;
+ }
+ auto & dp = dparams[seq_id];
+
+ const int32_t beg = i_block_beg[seq_id];
- const llama_token id = cur_p->data[0].id;
+ auto & result = *dp.result;
- if (cur_p->data[0].p < params.p_min) {
+ const float * lattice = llama_get_embeddings_nextn(ctx_dft);
+ GGML_ASSERT(lattice && "LiliCorr scorer produced no lattice");
+
+ int32_t predecessor = 0;
+ for (int32_t i = 1; i < params.n_max + 1; ++i) {
+ const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
+ const int32_t top_k = lilicorr_top_k;
+ const float * scores = row + top_k + (size_t) predecessor * top_k;
+
+ predecessor = (int32_t) std::distance(scores,
+ std::max_element(scores, scores + top_k));
+ if (params.p_min > 0.0f) {
+ // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
+ float sum = 0.0f;
+ for (int32_t k = 0; k < top_k; ++k) {
+ sum += std::exp(scores[k] - scores[predecessor]);
+ }
+ if (1.0f / sum < params.p_min) {
break;
}
-
- common_sampler_accept(smpl, id, true);
-
- result.push_back(id);
}
+ result.push_back((llama_token) row[predecessor]);
}
if (result.size() < (size_t) params.n_min) {
@@ -1517,10 +1788,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
}
}
}
-
- void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
- // noop
- }
};
struct common_speculative_impl_draft_mtp : public common_speculative_impl {
@@ -2535,18 +2802,20 @@ const char * common_speculative_all_types_str() {
std::string common_speculative_type_to_str(common_speculative_type type) {
switch (type) {
- case COMMON_SPECULATIVE_TYPE_NONE: return "none";
- case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: return "draft-simple";
- case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3";
- case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: return "draft-mtp";
- case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: return "draft-dflash";
- case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: return "draft-dspark";
- case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple";
- case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k";
- case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v";
- case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: return "ngram-mod";
- case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: return "ngram-cache";
- default: return "unknown";
+ case COMMON_SPECULATIVE_TYPE_NONE: return "none";
+ case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: return "draft-simple";
+ case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3";
+ case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: return "draft-mtp";
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: return "draft-dflash";
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: return "draft-dspark";
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2: return "draft-dflash2";
+ case COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR: return "draft-lilicorr";
+ case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple";
+ case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k";
+ case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v";
+ case COMMON_SPECULATIVE_TYPE_NGRAM_MOD: return "ngram-mod";
+ case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: return "ngram-cache";
+ default: return "unknown";
}
}
@@ -2604,10 +2873,8 @@ std::vector<common_speculative_type> common_speculative_types_from_gguf(const st
return {};
}
- // the Markov head distinguishes draft-dspark from draft-dflash
- const auto type = gguf_find_tensor(gguf_ctx.get(), "markov_w1.weight") >= 0
- ? COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK
- : COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH;
+ // draft-dflash runs every DFlash family draft, the head of the loaded draft picks the driver
+ const auto type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH;
SPC_INF("auto-detected speculative type '%s' from the draft model metadata\n", common_speculative_type_to_str(type).c_str());
@@ -2632,6 +2899,8 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
case COMMON_SPECULATIVE_TYPE_DRAFT_MTP:
case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH:
case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK:
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2:
+ case COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR:
n_max = std::max(n_max, std::max(0, spec->draft.n_max));
break;
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
@@ -2803,7 +3072,10 @@ common_params common_base_params_to_speculative(const common_params & params) {
const bool has_block_draft = std::any_of(
params.speculative.types.begin(), params.speculative.types.end(),
[](common_speculative_type t) {
- return t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK;
+ return t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH ||
+ t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK ||
+ t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2 ||
+ t == COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR;
});
if (has_block_draft) {
// per-seq output positions: DFlash decodes anchor + n_max masks (n_max + 1); DSpark n_max -> +1 covers both
@@ -2919,6 +3191,14 @@ common_speculative * common_speculative_init(common_params_speculative & params,
{
uint32_t enabled_configs = common_get_enabled_speculative_configs(params.types);
+ // the DFlash family types share the draft context, so one run uses at most one of them
+ const uint32_t dflash_configs = enabled_configs & (
+ (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) | (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) |
+ (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2) | (1u << COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR));
+ if (dflash_configs & (dflash_configs - 1)) {
+ throw std::runtime_error("draft-dflash, draft-dflash2, draft-dspark and draft-lilicorr cannot be combined");
+ }
+
auto add_config_if_enabled = [&](common_speculative_type type, bool available = true) {
if (available && (enabled_configs & (1u << type))) {
configs.emplace_back(type, params);
@@ -2926,7 +3206,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
};
// when adding a new type - update here the logic above
- static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 11);
+ static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 13);
// this list here defines the priority of the speculators
// the one with highest priority are listed first
@@ -2937,10 +3217,12 @@ common_speculative * common_speculative_init(common_params_speculative & params,
add_config_if_enabled(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE);
add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE);
- add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, params.draft.ctx_dft != nullptr);
- add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, params.draft.ctx_dft != nullptr);
- add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, params.draft.ctx_dft != nullptr);
- add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, params.draft.ctx_dft != nullptr);
+ add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, params.draft.ctx_dft != nullptr);
+ add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, params.draft.ctx_dft != nullptr);
+ add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, params.draft.ctx_dft != nullptr);
+ add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, params.draft.ctx_dft != nullptr);
+ add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2, params.draft.ctx_dft != nullptr);
+ add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR, params.draft.ctx_dft != nullptr);
}
std::vector<std::unique_ptr<common_speculative_impl>> impls = {};
@@ -2961,13 +3243,43 @@ common_speculative * common_speculative_init(common_params_speculative & params,
impls.push_back(std::make_unique<common_speculative_impl_draft_mtp>(config.params, n_seq));
break;
}
- case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: {
- impls.push_back(std::make_unique<common_speculative_impl_draft_dflash>(config.params, n_seq));
- break;
- }
- case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: {
- impls.push_back(std::make_unique<common_speculative_impl_draft_dflash>(
- config.params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK));
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH:
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK:
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2:
+ case COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR: {
+ common_speculative_type type_model = COMMON_SPECULATIVE_TYPE_NONE;
+ switch (llama_model_dflash_variant(llama_get_model(config.params.draft.ctx_dft))) {
+ case LLAMA_DFLASH_VARIANT_DFLASH: type_model = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH; break;
+ case LLAMA_DFLASH_VARIANT_DFLASH2: type_model = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2; break;
+ case LLAMA_DFLASH_VARIANT_DSPARK: type_model = COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; break;
+ case LLAMA_DFLASH_VARIANT_LILICORR: type_model = COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR; break;
+ case LLAMA_DFLASH_VARIANT_NONE: break;
+ }
+ if (type_model == COMMON_SPECULATIVE_TYPE_NONE) {
+ throw std::runtime_error(string_format("--spec-type %s requires a DFlash draft model",
+ common_speculative_type_to_str(config.type).c_str()));
+ }
+
+ // draft-dflash runs every DFlash family draft, the other types must match the head of the draft
+ if (config.type != COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH && config.type != type_model) {
+ throw std::runtime_error(string_format("--spec-type %s does not match the draft model: use %s",
+ common_speculative_type_to_str(config.type).c_str(), common_speculative_type_to_str(type_model).c_str()));
+ }
+
+ switch (type_model) {
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK:
+ impls.push_back(std::make_unique<common_speculative_impl_draft_dspark>(config.params, n_seq));
+ break;
+ case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH2:
+ impls.push_back(std::make_unique<common_speculative_impl_draft_dflash2>(config.params, n_seq));
+ break;
+ case COMMON_SPECULATIVE_TYPE_DRAFT_LILICORR:
+ impls.push_back(std::make_unique<common_speculative_impl_draft_lilicorr>(config.params, n_seq));
+ break;
+ default:
+ impls.push_back(std::make_unique<common_speculative_impl_draft_dflash>(config.params, n_seq));
+ break;
+ }
break;
}
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: {
@@ -3020,8 +3332,8 @@ common_speculative * common_speculative_init(common_params_speculative & params,
}
// the n-gram drafters ignore p_min, and DSpark uses it as the cutoff of its confidence head
- const bool uses_p_min = std::any_of(configs.begin(), configs.end(), [](const common_speculative_config & c) {
- switch (c.type) {
+ const bool uses_p_min = std::any_of(impls.begin(), impls.end(), [](const auto & impl) {
+ switch (impl->type) {
case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK:
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K:
diff --git a/common/speculative.h b/common/speculative.h
index 4262327cb..949114f35 100644
--- a/common/speculative.h
+++ b/common/speculative.h
@@ -73,7 +73,7 @@ struct common_speculative_draft_params {
// the generated draft from the last _draft() call
llama_tokens * result;
- // candidate distribution per drafted token; set it to make draft-simple, draft-mtp and draft-dflash (DFlash2 only) sample
+ // candidate distribution per drafted token; set it to make draft-simple, draft-mtp and draft-dflash2 sample
std::vector<std::vector<llama_token_data>> * result_q = nullptr;
// the target's temp and seed, read only when the drafter samples probabilistically
diff --git a/docs/speculative.md b/docs/speculative.md
index 9072eedd3..48e9d757e 100644
--- a/docs/speculative.md
+++ b/docs/speculative.md
@@ -74,13 +74,41 @@ llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DFlash.gguf \
`--spec-draft-n-max` is clamped to the draft model's trained block size.
+DFlash2, DSpark and LiliCorr drafts are DFlash backbones with a head on the block, and the head is read from the draft model.
+`draft-dflash` runs each of them with the driver of its head.
+The other spec types must match the head of the draft model: `draft-dspark` needs the Markov head, `draft-dflash2` the candidate selector and `draft-lilicorr` the candidate scorer. On a mismatch, speculative decoding is not initialized.
+Without `--spec-type`, the type is detected from the draft model.
+
See:
- #22105
+### DFlash2 (`draft-dflash2`)
+
+[DFlash2](https://inco.ai/blog/dflash2/) adds a local convolution to the DFlash layers and a _candidate selector_ on the block.
+
+DFlash2 drafts convert like DFlash drafts, with `--target-model-dir`:
+
+```bash
+python convert_hf_to_gguf.py <dflash2-draft-dir> \
+ --target-model-dir <target-model-dir> --outtype bf16 --outfile DFlash2.gguf
+
+llama-server -m target.gguf -md DFlash2.gguf \
+ --spec-type draft-dflash2 --spec-draft-n-max 7 -fa on --jinja
+```
+
+`--spec-draft-p-min P` truncates the draft at the first step whose selector probability falls below `P`.
+
+With `--spec-draft-sampling probabilistic`, the draft samples the path at the request's temperature
+and the target verifies it by rejection sampling.
+
+See:
+
+- #27816, #30663
+
### DSpark (`draft-dspark`)
-DSpark extends DFlash with a semi-autoregressive _Markov head_: the draft still emits a whole
+[DSpark](https://arxiv.org/pdf/2607.05147) extends DFlash with a semi-autoregressive _Markov head_: the draft still emits a whole
block per forward pass, but each block position's logits are biased by a low-rank term keyed on
the previous token, chained in-graph across the block. This keeps drafting at one decode per
block while recovering some of the left-to-right signal that pure block diffusion loses.
@@ -115,6 +143,26 @@ See:
- #25173
+### LiliCorr (`draft-lilicorr`)
+
+[LiliCorr](https://arxiv.org/pdf/2608.20530) puts a _candidate scorer_ on the DFlash block: a small transformer over the top-k candidates
+of each block position, conditioned on the target's hidden state at the last position, scores the paths through the block like the DFlash2 selector.
+
+LiliCorr drafts (`LiLiCorrDraftModel`) convert with `--target-model-dir`:
+
+```bash
+python convert_hf_to_gguf.py <lilicorr-draft-dir> \
+ --target-model-dir <target-model-dir> --outtype bf16 --outfile LiliCorr.gguf
+
+llama-server -m target.gguf -md LiliCorr.gguf --spec-type draft-lilicorr -fa on --jinja
+```
+
+`--spec-draft-p-min P` truncates the draft at the first step whose scorer probability falls below `P`.
+
+See:
+
+- #30506
+
### n-gram Cache (`ngram-cache`)
An n-gram is a sequence of n tokens. The n-gram cache implementation maintains statistics about short n-gram sequences.
@@ -225,7 +273,7 @@ Use exactly one of these options:
### General Speculative Parameters
```
---spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-dspark|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod]
+--spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-dflash2|draft-dspark|draft-lilicorr|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod]
comma-separated list of types of speculative decoding to use
(default: none)
(env: LLAMA_ARG_SPEC_TYPE)
@@ -368,8 +416,10 @@ Specifies a comma-separated list of speculative decoding types to use.
| `none` | No speculative decoding (default) |
| `draft-simple` | Use a simple draft model for speculation |
| `draft-eagle3` | Use an EAGLE-3 draft model that reads the target's hidden states |
-| `draft-dflash` | Use a DFlash block-diffusion draft model that emits a block per step |
+| `draft-dflash` | Use a DFlash block-diffusion draft model that emits a block per step (also DFlash2, DSpark and LiliCorr drafts, resolved from the model) |
+| `draft-dflash2` | Use a DFlash2 draft model (DFlash + local convolution + candidate selector) |
| `draft-dspark` | Use a DSpark draft model (DFlash backbone + semi-autoregressive Markov head) |
+| `draft-lilicorr` | Use a LiliCorr draft model (DFlash + candidate scorer conditioned on the target's hidden state) |
| `draft-mtp` | Use Multi Token Prediction (MTP) heads from the main model |
| `ngram-cache` | Use n-gram cache lookup |
| `ngram-simple` | Use simple n-gram pattern matching |
diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp
index e4a878b82..ec2c03176 100644
--- a/examples/speculative-simple/speculative-simple.cpp
+++ b/examples/speculative-simple/speculative-simple.cpp
@@ -116,7 +116,12 @@ int main(int argc, char ** argv) {
// init the speculator
const auto & params_spec = params.speculative;
- struct common_speculative * spec = common_speculative_init(params.speculative, 1);
+ struct common_speculative * spec = nullptr;
+ try {
+ spec = common_speculative_init(params.speculative, 1);
+ } catch (const std::exception & e) {
+ LOG_ERR("%s: %s\n", __func__, e.what());
+ }
if (spec == nullptr) {
LOG_ERR("%s", "failed to initialize speculative decoding\n");
diff --git a/include/llama.h b/include/llama.h
index 181d840a6..01d8fb88a 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -228,6 +228,15 @@ extern "C" {
LLAMA_CONTEXT_TYPE_MTP = 1,
};
+ // the head on the block of a DFlash draft model
+ enum llama_dflash_variant {
+ LLAMA_DFLASH_VARIANT_NONE = -1, // not a DFlash draft model
+ LLAMA_DFLASH_VARIANT_DFLASH = 0,
+ LLAMA_DFLASH_VARIANT_DFLASH2 = 1,
+ LLAMA_DFLASH_VARIANT_DSPARK = 2,
+ LLAMA_DFLASH_VARIANT_LILICORR = 3,
+ };
+
// TODO: simplify (https://github.com/ggml-org/llama.cpp/pull/9294#pullrequestreview-2286561979)
typedef struct llama_token_data {
llama_token id; // token id
@@ -667,6 +676,9 @@ extern "C" {
// NextN / speculative draft model queries
//
+ // the DFlash variant the model loader picked from the head of the draft
+ LLAMA_API enum llama_dflash_variant llama_model_dflash_variant(const struct llama_model * model);
+
// DFlash2 selector: number of candidate tokens the draft keeps per block position from its logits
// (the packed lattice also holds the top_k x top_k pairwise transition scores); 0 if the model has no selector head
LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model);
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 7c75ab55e..a5435cb63 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -3606,6 +3606,22 @@ int32_t llama_model_n_layer_nextn(const llama_model * model) {
return model->hparams.n_layer_nextn;
}
+llama_dflash_variant llama_model_dflash_variant(const llama_model * model) {
+ const auto * dflash = dynamic_cast<const llama_model_dflash *>(model);
+ if (dflash == nullptr) {
+ return LLAMA_DFLASH_VARIANT_NONE;
+ }
+
+ switch (dflash->variant) {
+ case llama_model_dflash::VARIANT_DFLASH: return LLAMA_DFLASH_VARIANT_DFLASH;
+ case llama_model_dflash::VARIANT_DFLASH2: return LLAMA_DFLASH_VARIANT_DFLASH2;
+ case llama_model_dflash::VARIANT_DSPARK: return LLAMA_DFLASH_VARIANT_DSPARK;
+ case llama_model_dflash::VARIANT_LILICORR: return LLAMA_DFLASH_VARIANT_LILICORR;
+ }
+
+ return LLAMA_DFLASH_VARIANT_NONE;
+}
+
int32_t llama_model_dflash_selector_top_k(const llama_model * model) {
return model->hparams.dflash_selector_top_k;
}
diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
index 13020bb32..ea1e862a9 100644
--- a/src/models/dflash.cpp
+++ b/src/models/dflash.cpp
@@ -33,18 +33,17 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE, hparams.dflash_conv_group_size, false);
ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK, hparams.dflash_selector_rank, false);
ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K, hparams.dflash_selector_top_k, false);
+
+ const bool has_markov = ml.get_tensor_meta("markov_w1.weight") != nullptr;
+ const bool has_selector = ml.get_tensor_meta("selector_hidden.weight") != nullptr;
const bool has_lilicorr = ml.get_key(LLM_KV_LILICORR_HIDDEN_SIZE, hparams.lilicorr_hidden_size, false);
- ml.get_key(LLM_KV_LILICORR_CANDIDATE_TOPK, hparams.lilicorr_candidate_topk, has_lilicorr);
- ml.get_key(LLM_KV_LILICORR_NUM_HEADS, hparams.lilicorr_num_heads, has_lilicorr);
- ml.get_key(LLM_KV_LILICORR_NUM_LAYERS, hparams.lilicorr_layers, has_lilicorr);
- ml.get_key(LLM_KV_LILICORR_VECTOR_EPS, hparams.lilicorr_vector_eps, has_lilicorr);
- ml.get_key(LLM_KV_LILICORR_LOGIT_SCALE, hparams.lilicorr_logit_scale, has_lilicorr);
- if (has_lilicorr && (hparams.dflash_block_size < 2 || hparams.lilicorr_hidden_size == 0 ||
- hparams.lilicorr_num_heads == 0 || hparams.lilicorr_hidden_size % hparams.lilicorr_num_heads ||
- hparams.lilicorr_candidate_topk == 0 || hparams.lilicorr_layers == 0 ||
- hparams.dflash_selector_top_k > 0 ||
- (uint64_t) hparams.lilicorr_candidate_topk * (hparams.lilicorr_candidate_topk + 1) > hparams.n_embd)) {
- throw std::runtime_error("invalid LiLiCorr DFlash metadata");
+ if (has_markov) {
+ variant = VARIANT_DSPARK;
+ } else if (has_selector) {
+ variant = VARIANT_DFLASH2;
+ } else if (has_lilicorr) {
+ variant = VARIANT_LILICORR;
+ load_hparams_lilicorr(ml);
}
if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) {
@@ -62,87 +61,16 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
}
LLAMA_LOG_INFO("%s: DFlash extract_layers = [%s]\n", __func__, layers.c_str());
- // DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);
if (hparams.dsv4_hc_mult > 0) {
- ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
- ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
- ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
- ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
- ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale);
- ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm);
- ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
- ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all);
- if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {
- hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;
- }
- ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count);
- ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank);
- ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
- ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
- ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false);
- ml.get_key(LLM_KV_HYPER_CONNECTION_SINGLE_PASS, hparams.dsv4_hc_single_pass, false);
-
- GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero
-
- if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
- throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring");
- }
- for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
- if (hparams.dsv4_compress_ratios[il] != 0) {
- throw std::runtime_error("DSpark DSV4 draft expects uncompressed attention on all stages");
- }
- }
-
- GGML_ASSERT(hparams.n_swa > 0);
- if (hparams.dsv4_hc_single_pass) {
- // DeepSeek-V4.1 DSpark: block queries see the main-KV ring plus the block via the mask row of the
- // first query, so that row's SWA window must reach window_size positions before the anchor
- hparams.n_swa += 1;
- }
- hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
- hparams.set_swa_pattern(0);
- for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
- hparams.is_swa_impl[il] = true;
- }
- hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
- hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
-
- type = LLM_TYPE_UNKNOWN;
- return;
- }
-
- // Mistral Large 4 DSpark: dense MLA stages, one sliding window, plain YaRN
- if (ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv, false)) {
- ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
- ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
- ml.get_key(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul, false);
- ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
-
- if (hparams.n_swa == 0) {
- throw std::runtime_error("DSpark Mistral Large 4 draft expects a sliding window");
- }
- if (hparams.n_embd_head_k_full != hparams.n_lora_kv + hparams.n_rot()) {
- throw std::runtime_error("DSpark Mistral Large 4 draft expects key_length == kv_lora_rank + rope dim");
- }
-
- // same as DSV4.1 DSpark: block uses the window of the anchor
- hparams.n_swa += 1;
- hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
- hparams.set_swa_pattern(0);
- for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
- hparams.is_swa_impl[il] = true;
- }
- hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
- hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
-
- type = LLM_TYPE_UNKNOWN;
- return;
- }
-
- // optional interleaved sliding-window attention with per-layer pattern array.
- // DFlash has a single rope, so the SWA rope == main rope.
- if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
+ load_hparams_dsv4(ml);
+ backbone = hparams.dsv4_hc_single_pass ? BACKBONE_DSV41 : BACKBONE_DSV4;
+ } else if (ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv, false)) {
+ load_hparams_mistral_large4(ml);
+ backbone = BACKBONE_MISTRAL_LARGE4;
+ } else if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
+ // optional interleaved sliding-window attention with per-layer pattern array.
+ // DFlash has a single rope, so the SWA rope == main rope.
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
@@ -168,109 +96,11 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
LLAMA_LOG_INFO("%s: DFlash using d2t mapping (draft_vocab_size = %lld)\n", __func__, (long long) n_vocab_draft);
}
- // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
- const struct ggml_tensor * markov_meta = ml->get_tensor_meta("markov_w1.weight");
- if (markov_meta) {
- const int64_t dspark_markov_rank = markov_meta->ne[0];
-
- dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);
- dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0);
- dspark_markov_w2_s = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
-
- dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED);
- dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED);
-
- dspark_joint_proj = create_tensor(tn(LLM_TENSOR_DSPARK_JOINT_PROJ, "weight"), { n_embd + 2*dspark_markov_rank, 3*dspark_markov_rank }, TENSOR_NOT_REQUIRED);
- dspark_joint_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_JOINT_PROJ, "bias"), { 3*dspark_markov_rank }, dspark_joint_proj ? 0 : TENSOR_NOT_REQUIRED);
-
- LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);
- }
-
- const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight");
- const bool conv_meta = hparams.dflash_conv_kernel_size > 0 && hparams.dflash_conv_group_size > 0;
- if (conv_meta && n_embd % hparams.dflash_conv_group_size != 0) {
- throw std::runtime_error("DFlash hidden size must be divisible by conv_group_size");
- }
- if (selector_meta) {
- const int64_t rank = hparams.dflash_selector_rank;
- if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 ||
- hparams.dflash_conv_kernel_size <= 0 || hparams.dflash_conv_group_size <= 0) {
- throw std::runtime_error("DFlash2 model is missing conv/selector metadata");
- }
- if (n_embd % hparams.dflash_conv_group_size != 0) {
- throw std::runtime_error("DFlash2 hidden size must be divisible by conv_group_size");
- }
- if (n_embd < hparams.dflash_selector_top_k * (hparams.dflash_selector_top_k + 1)) {
- throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice");
- }
-
- dflash_selector_prev = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV, "weight"), { rank, n_vocab }, 0);
- dflash_selector_next = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT, "weight"), { rank, n_vocab }, 0);
- dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0);
-
- LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__,
- hparams.dflash_conv_kernel_size, hparams.dflash_conv_group_size,
- hparams.dflash_selector_rank, hparams.dflash_selector_top_k);
- }
-
- if (hparams.lilicorr_layers > 0) {
- const int64_t hidden = hparams.lilicorr_hidden_size;
- const int64_t top_k = hparams.lilicorr_candidate_topk;
- const int64_t heads = hparams.lilicorr_num_heads;
- uint32_t factor = 0, mlp = 0;
- ml->get_key(LLM_KV_LILICORR_FACTOR_DIM, factor);
- ml->get_key(LLM_KV_LILICORR_MLP_HIDDEN_SIZE, mlp);
- const int64_t slots = hparams.dflash_block_size - 1;
- const int feature_bid = n_layer_all + hparams.lilicorr_layers;
- if (factor == 0 || mlp == 0 || hparams.lilicorr_layers > layers.size()) {
- throw std::runtime_error("invalid LiLiCorr head dimensions");
- }
- if (hidden != n_embd) {
- lilicorr_token_proj = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "weight"), { n_embd, hidden }, 0);
- lilicorr_token_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "bias"), { hidden }, 0);
- lilicorr_token_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
- }
- lilicorr_pass_hidden_proj = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "weight"), { hparams.n_embd, hidden }, 0);
- lilicorr_pass_hidden_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "bias"), { hidden }, 0);
- lilicorr_pass_hidden_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
- lilicorr_feature_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", feature_bid), { 5 }, 0);
- lilicorr_feature_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", feature_bid), { 5 }, 0);
- lilicorr_feature_mlp_1 = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", feature_bid), { 5, hidden }, 0);
- lilicorr_feature_mlp_1_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", feature_bid), { hidden }, 0);
- lilicorr_feature_mlp_3 = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", feature_bid), { hidden, hidden }, 0);
- lilicorr_feature_mlp_3_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", feature_bid), { hidden }, 0);
- lilicorr_slot_embedding = create_tensor(tn(LLM_TENSOR_LILICORR_SLOT_EMBD, "weight"), { hidden, slots }, 0);
- lilicorr_rank_embedding = create_tensor(tn(LLM_TENSOR_LILICORR_RANK_EMBD, "weight"), { hidden, top_k }, 0);
- lilicorr_attention_bias = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_BIAS, "bias"), { slots * top_k, slots * top_k, heads }, 0);
- lilicorr_context_proj = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "weight"), { n_embd, hidden }, 0);
- lilicorr_context_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "bias"), { hidden }, 0);
- lilicorr_context_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
- for (uint32_t i = 0; i < hparams.lilicorr_layers; ++i) {
- const int bid = n_layer_all + i;
- layers[i].lilicorr_attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", bid), { hidden }, 0);
- layers[i].lilicorr_mlp_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", bid), { hidden }, 0);
- layers[i].lilicorr_attn_in_proj = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "weight", i), { hidden, 3 * hidden }, 0);
- layers[i].lilicorr_attn_in_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "bias", i), { 3 * hidden }, 0);
- layers[i].lilicorr_attn_out_proj = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", bid), { hidden, hidden }, 0);
- layers[i].lilicorr_attn_out_proj_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", bid), { hidden }, 0);
- layers[i].lilicorr_attn_out_proj_s = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
- layers[i].lilicorr_mlp_0 = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", bid), { hidden, mlp }, 0);
- layers[i].lilicorr_mlp_0_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", bid), { mlp }, 0);
- layers[i].lilicorr_mlp_0_s = create_tensor(tn(LLM_TENSOR_FFN_UP, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
- layers[i].lilicorr_mlp_2 = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", bid), { mlp, hidden }, 0);
- layers[i].lilicorr_mlp_2_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", bid), { hidden }, 0);
- layers[i].lilicorr_mlp_2_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
- }
- lilicorr_output_norm = create_tensor(tn(LLM_TENSOR_LILICORR_OUTPUT_NORM, "weight"), { hidden }, 0);
- lilicorr_anchor_norm = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_NORM, "weight"), { hidden }, 0);
- lilicorr_factor_input_proj = create_tensor(tn(LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "weight"), { 3 * hidden, hidden }, 0);
- lilicorr_factor_input_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "bias"), { hidden }, 0);
- lilicorr_out_head = create_tensor(tn(LLM_TENSOR_LILICORR_OUT_HEAD, "weight"), { hidden, factor }, 0);
- lilicorr_out_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_OUT_HEAD, "bias"), { factor }, 0);
- lilicorr_in_head = create_tensor(tn(LLM_TENSOR_LILICORR_IN_HEAD, "weight"), { hidden, factor }, 0);
- lilicorr_in_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_IN_HEAD, "bias"), { factor }, 0);
- lilicorr_anchor_out_head = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "weight"), { hidden, factor }, 0);
- lilicorr_anchor_out_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "bias"), { factor }, 0);
+ switch (variant) {
+ case VARIANT_DSPARK: load_tensors_dspark(n_vocab_draft); break;
+ case VARIANT_DFLASH2: load_tensors_dflash2(); break;
+ case VARIANT_LILICORR: load_tensors_lilicorr(); break;
+ case VARIANT_DFLASH: break;
}
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
@@ -283,85 +113,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
// a draft with its own embeddings + head references no target tensors and can run on devices the target does not use (e.g. -devd with a tensor-split target)
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED);
- if (hparams.n_lora_kv > 0) {
- const int64_t kv_lora_rank = hparams.n_lora_kv;
- const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
- const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
- const int64_t n_embd_head_nope = n_embd_head_k_mla - n_rot;
-
- dspark_mask_embd = create_tensor(tn(LLM_TENSOR_DSPARK_MASK_EMBD, "weight"), { n_embd, hparams.dflash_block_size - 1 }, 0);
-
- for (int i = 0; i < n_layer; ++i) {
- auto & layer = layers[i];
-
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
- layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_head * n_embd_head_k_mla }, 0);
- layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), { n_embd, kv_lora_rank + n_rot }, 0);
- layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), { kv_lora_rank }, 0);
- layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), { n_embd_head_nope, kv_lora_rank, n_head }, 0);
- layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), { kv_lora_rank, n_embd_head_v_mla, n_head }, 0);
- layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_head * n_embd_head_v_mla, n_embd }, 0);
-
- layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0);
- layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);
- layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);
- layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0);
- }
- return;
+ switch (backbone) {
+ case BACKBONE_BASE: load_tensors_layers_base(); break;
+ case BACKBONE_DSV4:
+ case BACKBONE_DSV41: load_tensors_layers_dsv4(); break;
+ case BACKBONE_MISTRAL_LARGE4: load_tensors_layers_mistral_large4(); break;
}
+}
- if (hparams.dsv4_hc_mult > 0) {
- const int64_t q_lora_rank = hparams.n_lora_q;
- const int64_t n_ff_exp = hparams.n_ff_exp();
- const int64_t n_expert_shared = hparams.n_expert_shared;
- const int64_t n_embd_head = hparams.n_embd_head_k();
- const int64_t o_groups = hparams.dsv4_o_group_count;
- const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
- const int64_t hc_mult = hparams.dsv4_hc_mult;
- const int64_t hc_dim = hc_mult * n_embd;
- const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
-
- // single-pass mHC: collapse with the last FFN pre-mix, no hc_head
- if (!hparams.dsv4_hc_single_pass) {
- hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0);
- hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0);
- hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);
- }
-
- for (int i = 0; i < n_layer; ++i) {
- auto & layer = layers[i];
-
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
- layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);
- layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
- layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);
- layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);
- layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);
- layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);
- layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, TENSOR_ALLOW_RESHAPE);
- layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
-
- layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
- layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);
- layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
- layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
- layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);
- layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);
-
- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
- layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
- layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
-
- layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
- layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
-
- layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
- layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0);
- layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
- }
- return;
- }
+void llama_model_dflash::load_tensors_layers_base() {
+ LLAMA_LOAD_LOCALS;
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
@@ -388,16 +149,6 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0);
-
- if (conv_meta) {
- const int64_t kernel = hparams.dflash_conv_kernel_size;
- const int64_t groups = n_embd / hparams.dflash_conv_group_size;
- const int64_t projected = 2 * kernel * groups;
- layer.dflash_attn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);
- layer.dflash_attn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);
- layer.dflash_ffn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);
- layer.dflash_ffn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);
- }
}
}
@@ -451,612 +202,153 @@ static ggml_tensor * build_inp_target_features(llm_graph_context & g, llm_graph_
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;
- auto & res = g.res;
-
- ggml_tensor * w1 = model.dspark_markov_w1;
- ggml_tensor * w2 = model.dspark_markov_w2;
- GGML_ASSERT(w1 && w2 && "DSpark markov weights not loaded");
-
- // confidence head is optional
- const bool has_conf = model.dspark_conf_proj != nullptr;
-
- // rnn head: gated state carried across block positions
- const bool is_rnn = model.dspark_joint_proj != nullptr;
- GGML_ASSERT(!is_rnn || hidden != nullptr);
-
- ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens]
- const int64_t n_vocab = base->ne[0];
- const int64_t n_tok = base->ne[1];
-
- const auto it = model.gguf_kv.find("dflash.block_size");
- GGML_ASSERT(it != model.gguf_kv.end() && "DSpark draft requires 'dflash.block_size' in GGUF metadata");
- const int64_t block_size = std::stoi(it->second);
- GGML_ASSERT(block_size > 0);
-
- // bonus anchor (SpecForge exports): slot 0 is a bonus token, not a prediction slot
- const auto it_anchor = model.gguf_kv.find("dflash.sample_from_anchor");
- const bool sample_from_anchor = it_anchor == model.gguf_kv.end() || it_anchor->second == "true";
- const int64_t i_draft_beg = sample_from_anchor ? 0 : 1;
-
- const int64_t n_blocks = g.ubatch.n_seqs_unq;
- GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks");
- // runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size
- const int64_t block_drafts = n_tok / n_blocks;
- if (block_drafts > block_size) {
- return;
- }
-
- // anchor (committed last) token of every block: token 0 of each block, i.e. a strided view
- const size_t token_stride = (size_t) block_drafts * tokens->nb[0];
- const size_t base_stride = (size_t) block_drafts * base->nb[1];
-
- ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0);
- prev = ggml_cont_1d(ctx0, prev, n_blocks);
-
- ggml_tensor * cat = nullptr;
- ggml_tensor * cat_conf = nullptr;
-
- if (!sample_from_anchor) {
- // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column
- cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));
- if (has_conf) {
- cat_conf = ggml_sigmoid(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0));
- }
- }
-
- ggml_tensor * state = nullptr; // [R, n_blocks]
-
- // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final
- // token pick, not the Markov conditioning path
- for (int64_t i = i_draft_beg; i < block_drafts; ++i) {
- ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks]
- ggml_tensor * markov = w1_prev;
-
- if (is_rnn) {
- const int64_t rank = w1_prev->ne[0];
- if (state == nullptr) {
- state = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, rank, n_blocks), 0.0f);
- }
+// injection batch: the target features fused through the encoder
+ggml_tensor * llama_model_dflash::build_inject_inp(llm_graph_context & g, const llama_model & model) {
+ const int64_t n_embd_inp = g.hparams.n_embd_inp_enc();
- ggml_tensor * hidden_i = ggml_view_2d(ctx0, hidden, hidden->ne[0], n_blocks,
- (size_t) block_drafts * hidden->nb[1], i*hidden->nb[1]);
- ggml_tensor * joint = ggml_concat(ctx0, ggml_concat(ctx0, state, w1_prev, 0), ggml_cont(ctx0, hidden_i), 0);
- joint = ggml_add(ctx0, ggml_mul_mat(ctx0, model.dspark_joint_proj, joint), model.dspark_joint_proj_b); // [3R, n_blocks]
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
- const auto chunk = [&](int64_t k) {
- return ggml_cont(ctx0, ggml_view_2d(ctx0, joint, rank, n_blocks, joint->nb[1], k*rank*ggml_element_size(joint)));
- };
- ggml_tensor * gate = ggml_sigmoid(ctx0, chunk(0));
- ggml_tensor * cand = ggml_tanh(ctx0, chunk(1));
+ ggml_tensor * inp_target = build_inp_target_features(g, *inp, n_embd_inp);
+ g.cb(inp_target, "inp_target_features", -1);
- // state = gate * state + (1 - gate) * cand
- state = ggml_add(ctx0, cand, ggml_mul(ctx0, gate, ggml_sub(ctx0, state, cand)));
- markov = ggml_tanh(ctx0, chunk(2));
- }
-
- ggml_tensor * bias = g.build_lora_mm(w2, markov, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks]
- if (model.d2t) {
- // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others)
- const int64_t n_draft_vocab = bias->ne[0];
- ggml_tensor * full = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_blocks), 0.0f);
- bias = ggml_set_rows(ctx0, full,
- ggml_reshape_3d(ctx0, bias, 1, n_draft_vocab, n_blocks),
- ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1));
- bias = ggml_reshape_2d(ctx0, bias, n_vocab, n_blocks);
- }
-
- // position i of every block: strided view [n_vocab, n_blocks]
- ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]);
- ggml_tensor * col = ggml_add(ctx0, base_i, bias);
-
- cat = cat ? ggml_concat(ctx0, cat, col, 1) : col;
-
- if (has_conf) {
- // confidence head input: predicts per-position acceptance
- ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok]
- // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks]
- ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,
- (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);
- ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), markov, 0);
- ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);
- if (model.dspark_conf_proj_b) {
- conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);
- }
- conf = ggml_sigmoid(ctx0, conf);
+ g.res->add_input(std::move(inp));
- cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;
- }
-
- if (i + 1 < block_drafts) {
- prev = ggml_argmax(ctx0, col);
- }
- }
-
- // cat is position-major; restore ubatch block-major order
- ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, block_drafts);
- out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks]
- out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok);
-
- if (has_conf) {
- ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts);
- conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3));
- conf = ggml_reshape_2d(ctx0, conf, 1, n_tok);
-
- // note: broadcast the [1, n_tok] confidences to n_embd-wide rows to be able to reuse `llama_get_embeddings_nextn`
- conf = ggml_repeat(ctx0, conf, res->t_embd);
- res->t_h_nextn = conf;
- ggml_build_forward_expand(g.gf, conf);
- }
-
- res->t_logits = out;
- ggml_build_forward_expand(g.gf, out);
+ ggml_tensor * inp_g = g.build_lora_mm(model.fc, inp_target, model.fc_s);
+ return g.build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
}
-static ggml_tensor * build_dflash2_conv(
- llm_graph_context & g,
- ggml_tensor * hidden,
- ggml_tensor * dynamic,
- ggml_tensor * base,
- int side) {
- const auto & hparams = g.hparams;
- const int64_t hidden_size = hidden->ne[0];
- const int64_t n_tokens = hidden->ne[1];
- const int64_t n_blocks = g.ubatch.n_seqs_unq;
- const int64_t kernel_size = hparams.dflash_conv_kernel_size;
- const int64_t group_size = hparams.dflash_conv_group_size;
- const int64_t n_groups = hidden_size / group_size;
-
- GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
- GGML_ASSERT(dynamic && base && side >= 0 && side < 2);
-
- const int64_t block_size = n_tokens / n_blocks;
- ggml_context * ctx0 = g.ctx0;
- // ggml_cont copies even when the tensor is already contiguous
- if (!ggml_is_contiguous(hidden) || hidden->ne[1] != n_tokens) {
- hidden = ggml_cont_2d(ctx0, hidden, hidden_size, n_tokens);
- }
- if (!ggml_is_contiguous(dynamic) || dynamic->ne[1] != n_tokens) {
- dynamic = ggml_cont_2d(ctx0, dynamic, dynamic->ne[0], n_tokens);
+// tok_embd from the target model (shared via ctx_other) when the draft has none
+ggml_tensor * llama_model_dflash::get_tok_embd(const llm_graph_context & g, const llama_model & model) {
+ if (model.tok_embd) {
+ return model.tok_embd;
}
- ggml_tensor * blocks = ggml_reshape_3d(ctx0, hidden, hidden_size, block_size, n_blocks);
- ggml_tensor * coeffs = ggml_reshape_4d(ctx0, dynamic, n_groups, kernel_size, 2, n_tokens);
- ggml_tensor * coeffs_side = ggml_view_3d(ctx0, coeffs, n_groups, kernel_size, n_tokens,
- coeffs->nb[1], coeffs->nb[3], side * coeffs->nb[2]);
-
- ggml_tensor * coeff_all = ggml_cont(ctx0, coeffs_side);
- coeff_all = ggml_reshape_4d(ctx0, coeff_all, 1, n_groups, kernel_size, n_tokens);
- coeff_all = ggml_repeat_4d(ctx0, coeff_all, group_size, n_groups, kernel_size, n_tokens);
-
- ggml_tensor * base_side = ggml_reshape_4d(ctx0,
- ggml_view_1d(ctx0, base, hidden_size * kernel_size, side * base->nb[2]),
- group_size, n_groups, kernel_size, 1);
-
- ggml_tensor * weight_all = ggml_add(ctx0, coeff_all, base_side);
-
- // taps at or past block_size only read the left padding and add nothing
- const int64_t n_taps = std::min(kernel_size, block_size);
-
- ggml_tensor * result = nullptr;
- for (int64_t tap = 0; tap < n_taps; ++tap) {
- ggml_tensor * values = blocks;
- if (tap > 0) {
- ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks,
- blocks->nb[1], blocks->nb[2], 0);
- values = ggml_pad_ext(ctx0, previous, 0, 0, tap, 0, 0, 0, 0, 0);
- }
- values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens);
-
- ggml_tensor * weight = ggml_reshape_2d(ctx0,
- ggml_cont(ctx0, ggml_view_4d(ctx0, weight_all, group_size, n_groups, 1, n_tokens,
- weight_all->nb[1], weight_all->nb[2], weight_all->nb[3], tap * weight_all->nb[2])),
- hidden_size, n_tokens);
+ GGML_ASSERT(g.cparams.ctx_other != nullptr);
+ const auto * model_other = llama_get_model(g.cparams.ctx_other);
- ggml_tensor * term = ggml_mul(ctx0, weight, values);
- result = result ? ggml_add(ctx0, result, term) : term;
- }
- return result;
+ GGML_ASSERT(model_other->tok_embd != nullptr && "DFlash draft requires the target model's token embeddings");
+ return model_other->tok_embd;
}
-// DFlash2 selector: top-k candidates per block position plus the pairwise
-// transition scores, packed into the nextn output slot for the CPU-side walk.
-static void build_dflash2_selector(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {
- ggml_context * ctx0 = g.ctx0;
- auto & res = g.res;
-
- const auto & hparams = g.hparams;
- const int64_t n_tokens = g.n_tokens;
- const int64_t n_embd = g.n_embd;
-
- const int64_t top_k = hparams.dflash_selector_top_k;
- const int64_t rank = hparams.dflash_selector_rank;
- const int64_t n_blocks = g.ubatch.n_seqs_unq;
- GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
- GGML_ASSERT(res->t_logits->ne[1] == n_tokens);
- if (!tokens) {
- return;
+// lm_head and its scale from the target model (shared via ctx_other) when the draft has none
+std::pair<ggml_tensor *, ggml_tensor *> llama_model_dflash::get_lm_head(const llm_graph_context & g, const llama_model & model) {
+ if (model.output) {
+ return { model.output, model.output_s };
}
+ GGML_ASSERT(g.cparams.ctx_other != nullptr);
+ const auto * model_other = llama_get_model(g.cparams.ctx_other);
- const int64_t tokens_per_block = n_tokens / n_blocks;
- const int64_t block_size = std::min<int64_t>(tokens_per_block, hparams.dflash_block_size);
- const int64_t row_used = top_k + top_k * top_k;
-
- ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k);
- ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens);
- ggml_tensor * unary = ggml_reshape_2d(ctx0,
- ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens);
- ggml_tensor * gate = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd);
-
- // Everything below indexes [.., tokens_per_block, n_blocks]: the block
- // position varies fastest, sequences are the outer dimension.
- ggml_tensor * cand_blk = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks);
- ggml_tensor * unary_blk = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks);
- ggml_tensor * gate_blk = ggml_reshape_3d(ctx0, gate, rank, tokens_per_block, n_blocks);
-
- // a position's score reads only the candidate sets at pos-1 and pos, so a run
- // of positions has no internal dependency and scores in one batched matmul
- auto score_run = [&](int64_t beg_pos, int64_t n_pos, ggml_tensor * pred_ids) {
- ggml_tensor * cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, n_pos, n_blocks,
- cand_blk->nb[1], cand_blk->nb[2], beg_pos * cand_blk->nb[1]));
- ggml_tensor * unary_run = ggml_cont(ctx0, ggml_view_3d(ctx0, unary_blk, top_k, n_pos, n_blocks,
- unary_blk->nb[1], unary_blk->nb[2], beg_pos * unary_blk->nb[1]));
- ggml_tensor * gate_run = ggml_cont(ctx0, ggml_view_3d(ctx0, gate_blk, rank, n_pos, n_blocks,
- gate_blk->nb[1], gate_blk->nb[2], beg_pos * gate_blk->nb[1]));
-
- const int64_t n_pred = pred_ids->ne[0] / (n_pos * n_blocks);
-
- ggml_tensor * successor = ggml_reshape_4d(ctx0,
- ggml_get_rows(ctx0, model.dflash_selector_next, ggml_reshape_1d(ctx0, cand_run, top_k * n_pos * n_blocks)),
- rank, top_k, n_pos, n_blocks);
- ggml_tensor * predecessor = ggml_reshape_4d(ctx0,
- ggml_get_rows(ctx0, model.dflash_selector_prev, pred_ids),
- rank, n_pred, n_pos, n_blocks);
-
- ggml_tensor * gate_bcast = ggml_reshape_4d(ctx0, gate_run, rank, 1, n_pos, n_blocks);
- ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, gate_bcast, predecessor));
- ggml_tensor * score = ggml_mul_mat(ctx0, successor, cond);
- if (n_pred == 1) {
- score = ggml_repeat_4d(ctx0, score, top_k, top_k, n_pos, n_blocks);
- }
- ggml_tensor * unary_bcast = ggml_reshape_4d(ctx0, unary_run, top_k, 1, n_pos, n_blocks);
- score = ggml_add(ctx0, score, ggml_repeat(ctx0, unary_bcast, score));
-
- ggml_tensor * row = ggml_concat(ctx0,
- ggml_cast(ctx0, cand_run, GGML_TYPE_F32),
- ggml_reshape_3d(ctx0, score, top_k * top_k, n_pos, n_blocks), 0);
- return ggml_pad(ctx0, row, n_embd - row_used, 0, 0, 0);
- };
-
- ggml_tensor * packed = ggml_fill(ctx0,
- ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f);
-
- if (block_size > 1) {
- // Position 1 alone: its predecessor is the anchor token, one id per
- // sequence rather than a candidate set.
- ggml_tensor * anchor_ids = ggml_cont_1d(ctx0,
- ggml_view_2d(ctx0, tokens, 1, n_blocks, tokens_per_block * tokens->nb[0], 0), n_blocks);
- packed = ggml_concat(ctx0, packed, score_run(1, 1, anchor_ids), 1);
- }
- if (block_size > 2) {
- ggml_tensor * prev_ids = ggml_reshape_1d(ctx0,
- ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, block_size - 2, n_blocks,
- cand_blk->nb[1], cand_blk->nb[2], cand_blk->nb[1])),
- top_k * (block_size - 2) * n_blocks);
- packed = ggml_concat(ctx0, packed, score_run(2, block_size - 2, prev_ids), 1);
- }
-
- packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks);
- g.cb(packed, "dflash2_lattice", -1);
- res->t_h_nextn = packed;
- ggml_build_forward_expand(g.gf, packed);
-}
-
-static void build_lilicorr_scores(llm_graph_context & g, const llama_model & model,
- ggml_tensor * tok_embd, ggml_tensor * anchors) {
- ggml_context * ctx0 = g.ctx0;
- auto & res = g.res;
-
- const auto & hparams = g.hparams;
- const int64_t n_tokens = g.n_tokens;
- const int64_t n_embd = g.n_embd;
-
- const int64_t hidden = hparams.lilicorr_hidden_size;
- const int64_t top_k = hparams.lilicorr_candidate_topk;
- const int64_t heads = hparams.lilicorr_num_heads;
- const int64_t n_blocks = g.ubatch.n_seqs_unq;
- GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
- GGML_ASSERT(res->t_logits->ne[1] == n_tokens);
- GGML_ASSERT(anchors && tok_embd);
-
- const int64_t tokens_per_block = n_tokens / n_blocks;
- const int64_t block_size = hparams.dflash_block_size;
- const int64_t row_used = top_k + top_k * top_k;
- const int64_t slots = block_size - 1;
- const int64_t nodes = slots * top_k;
- const int64_t factor_dim = model.lilicorr_out_head->ne[1];
- GGML_ASSERT(tokens_per_block == block_size);
-
- ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k);
- ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens);
- ggml_tensor * candidate_logits = ggml_reshape_2d(ctx0,
- ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens);
- // Rank features and embeddings require candidates in descending logit order.
- ggml_tensor * order = ggml_argsort(ctx0, candidate_logits, GGML_SORT_ORDER_DESC);
- ggml_tensor * candidate_rows = ggml_reshape_3d(ctx0, candidates, 1, top_k, n_tokens);
- candidates = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, candidate_rows, order), top_k, n_tokens);
- ggml_tensor * prob_rows = ggml_reshape_3d(ctx0, ggml_soft_max(ctx0, res->t_logits), 1, res->t_logits->ne[0], n_tokens);
- ggml_tensor * probs = ggml_reshape_2d(ctx0,
- ggml_get_rows(ctx0, prob_rows, candidates), top_k, n_tokens);
-
- ggml_tensor * cand_run = ggml_reshape_3d(ctx0, candidates, top_k, n_tokens / n_blocks, n_blocks);
- ggml_tensor * probs_run = ggml_reshape_3d(ctx0, probs, top_k, n_tokens / n_blocks, n_blocks);
- cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_run, top_k, slots, n_blocks,
- cand_run->nb[1], cand_run->nb[2], cand_run->nb[1]));
- probs_run = ggml_cont(ctx0, ggml_view_3d(ctx0, probs_run, top_k, slots, n_blocks,
- probs_run->nb[1], probs_run->nb[2], probs_run->nb[1]));
-
- ggml_tensor * lp = ggml_log(ctx0, probs_run);
- ggml_tensor * lp_blk = ggml_reshape_3d(ctx0, lp, top_k, slots, n_blocks);
- ggml_tensor * top_lp = ggml_view_3d(ctx0, lp_blk, 1, slots, n_blocks,
- lp_blk->nb[1], lp_blk->nb[2], 0);
- ggml_tensor * relative = ggml_sub(ctx0, lp_blk, top_lp);
- ggml_tensor * feature = ggml_reshape_4d(ctx0, lp_blk, 1, top_k, slots, n_blocks);
-
- ggml_tensor * rank = ggml_arange(ctx0, 0, top_k, 1);
- rank = ggml_repeat(ctx0, ggml_reshape_3d(ctx0, rank, top_k, 1, 1), lp_blk);
- ggml_tensor * first = ggml_scale_bias(ctx0, ggml_step(ctx0, rank), -1.0f, 1.0f);
-
- for (ggml_tensor * item : { probs_run, relative,
- ggml_scale(ctx0, rank, top_k == 1 ? 0.0f : 1.0f / (top_k - 1)), first }) {
- feature = ggml_concat(ctx0, feature, ggml_reshape_4d(ctx0, item, 1, top_k, slots, n_blocks), 0);
- }
- feature = ggml_reshape_2d(ctx0, feature, 5, nodes * n_blocks);
-
- ggml_tensor * embeds = ggml_get_rows(ctx0, tok_embd, ggml_reshape_1d(ctx0, cand_run, nodes * n_blocks));
- ggml_tensor * token_states = n_embd == hidden ? embeds : ggml_add(ctx0, g.build_lora_mm(model.lilicorr_token_proj, embeds, model.lilicorr_token_proj_s), model.lilicorr_token_proj_b);
-
- ggml_tensor * pass_blk = ggml_reshape_3d(ctx0, res->t_embd, n_embd, block_size, n_blocks);
- ggml_tensor * passes = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_view_3d(ctx0, pass_blk, n_embd, slots, n_blocks,
- pass_blk->nb[1], pass_blk->nb[2], pass_blk->nb[1])),
- n_embd, slots * n_blocks);
-
- ggml_tensor * pass_states = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_pass_hidden_proj, passes, model.lilicorr_pass_hidden_proj_s), model.lilicorr_pass_hidden_proj_b);
-
- pass_states = ggml_reshape_4d(ctx0, pass_states, hidden, 1, slots, n_blocks);
- pass_states = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, pass_states, hidden, top_k, slots, n_blocks), hidden, nodes * n_blocks);
-
- ggml_tensor * feature_states = ggml_norm(ctx0, feature, 1.0e-5f);
- feature_states = ggml_add(ctx0, ggml_mul(ctx0, feature_states, model.lilicorr_feature_norm), model.lilicorr_feature_norm_b);
- feature_states = ggml_silu(ctx0, ggml_add(ctx0, g.build_lora_mm(model.lilicorr_feature_mlp_1, feature_states), model.lilicorr_feature_mlp_1_b));
- feature_states = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_feature_mlp_3, feature_states), model.lilicorr_feature_mlp_3_b);
-
- ggml_tensor * slot = ggml_reshape_4d(ctx0, model.lilicorr_slot_embedding, hidden, 1, slots, 1);
- ggml_tensor * rank_embd = ggml_reshape_4d(ctx0, model.lilicorr_rank_embedding, hidden, top_k, 1, 1);
- ggml_tensor * pos = ggml_add(ctx0, ggml_repeat_4d(ctx0, slot, hidden, top_k, slots, n_blocks), ggml_repeat_4d(ctx0, rank_embd, hidden, top_k, slots, n_blocks));
- pos = ggml_reshape_2d(ctx0, pos, hidden, nodes * n_blocks);
-
- ggml_tensor * x = ggml_add(ctx0, ggml_add(ctx0, token_states, pass_states), feature_states);
- x = ggml_add(ctx0, x, pos);
- ggml_tensor * anchor_blk = ggml_reshape_3d(ctx0, anchors, n_embd, block_size, n_blocks);
- ggml_tensor * anchor = ggml_reshape_2d(ctx0,
- ggml_cont(ctx0, ggml_view_3d(ctx0, anchor_blk, n_embd, 1, n_blocks,
- anchor_blk->nb[1], anchor_blk->nb[2], 0)), n_embd, n_blocks);
- ggml_tensor * anchor_state = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_context_proj, anchor, model.lilicorr_context_proj_s), model.lilicorr_context_proj_b);
-
- const int64_t head_dim = hidden / heads;
- for (uint32_t layer = 0; layer < hparams.lilicorr_layers; ++layer) {
- const auto & weights = model.layers[layer];
- ggml_tensor * attn_input = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), weights.lilicorr_attn_norm);
- ggml_tensor * qkv = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_attn_in_proj, attn_input), weights.lilicorr_attn_in_proj_b);
- qkv = ggml_reshape_3d(ctx0, qkv, 3 * hidden, nodes, n_blocks);
- ggml_tensor * q = ggml_cont(ctx0, ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
- qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], 0));
- ggml_tensor * k = ggml_cont(ctx0, ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
- qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], hidden * qkv->nb[0]));
- ggml_tensor * v = ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
- qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], 2 * hidden * qkv->nb[0]);
- ggml_tensor * attn = ggml_soft_max_ext(ctx0, ggml_mul_mat(ctx0, k, q), model.lilicorr_attention_bias,
- 1.0f / std::sqrt(static_cast<float>(head_dim)), 0.0f);
- ggml_tensor * joined = ggml_out_prod(ctx0, v, ggml_transpose(ctx0, attn));
- joined = ggml_permute(ctx0, joined, 0, 2, 1, 3);
- ggml_tensor * attn_output = ggml_cont_2d(ctx0, joined, hidden, nodes * n_blocks);
- attn_output = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_attn_out_proj, attn_output, weights.lilicorr_attn_out_proj_s), weights.lilicorr_attn_out_proj_b);
- x = ggml_add(ctx0, x, attn_output);
- ggml_tensor * mlp = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), weights.lilicorr_mlp_norm);
- mlp = ggml_silu(ctx0, ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_mlp_0, mlp, weights.lilicorr_mlp_0_s), weights.lilicorr_mlp_0_b));
- mlp = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_mlp_2, mlp, weights.lilicorr_mlp_2_s), weights.lilicorr_mlp_2_b);
- x = ggml_add(ctx0, x, mlp);
- }
-
- x = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), model.lilicorr_output_norm);
- anchor_state = ggml_mul(ctx0, ggml_rms_norm(ctx0, anchor_state, hparams.f_norm_rms_eps), model.lilicorr_anchor_norm);
- ggml_tensor * anchor_nodes = ggml_reshape_3d(ctx0, anchor_state, hidden, 1, n_blocks);
- anchor_nodes = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, anchor_nodes, hidden, nodes, n_blocks, 1), hidden, nodes * n_blocks);
- ggml_tensor * factor_in = ggml_concat(ctx0, x, anchor_nodes, 0);
- factor_in = ggml_concat(ctx0, factor_in, ggml_mul(ctx0, x, anchor_nodes), 0);
- ggml_tensor * factors = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_factor_input_proj, factor_in), model.lilicorr_factor_input_proj_b);
- factors = ggml_silu(ctx0, factors);
-
- ggml_tensor * out = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_out_head, factors), model.lilicorr_out_head_b);
- out = ggml_reshape_4d(ctx0, ggml_l2_norm(ctx0, out, hparams.lilicorr_vector_eps), factor_dim, top_k, slots, n_blocks);
-
- ggml_tensor * in = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_in_head, factors), model.lilicorr_in_head_b);
- in = ggml_reshape_4d(ctx0, ggml_l2_norm(ctx0, in, hparams.lilicorr_vector_eps), factor_dim, top_k, slots, n_blocks);
-
- ggml_tensor * anchor_out = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_anchor_out_head, anchor_state), model.lilicorr_anchor_out_head_b);
- anchor_out = ggml_l2_norm(ctx0, anchor_out, hparams.lilicorr_vector_eps);
- anchor_out = ggml_reshape_3d(ctx0, anchor_out, factor_dim, 1, n_blocks);
-
- ggml_tensor * first_in = ggml_cont(ctx0, ggml_view_4d(ctx0, in, factor_dim, top_k, 1, n_blocks,
- in->nb[1], in->nb[2], in->nb[3], 0));
- ggml_tensor * first_score = ggml_scale(ctx0,
- ggml_mul_mat(ctx0, first_in, ggml_reshape_4d(ctx0, anchor_out, factor_dim, 1, 1, n_blocks)),
- hparams.lilicorr_logit_scale);
- ggml_tensor * scores = ggml_pad(ctx0,
- ggml_reshape_3d(ctx0, first_score, top_k, 1, n_blocks), top_k * top_k - top_k, 0, 0, 0);
- if (slots > 1) {
- ggml_tensor * next = ggml_cont(ctx0, ggml_view_4d(ctx0, in, factor_dim, top_k, slots - 1, n_blocks,
- in->nb[1], in->nb[2], in->nb[3], in->nb[2]));
- ggml_tensor * previous = ggml_cont(ctx0, ggml_view_4d(ctx0, out, factor_dim, top_k, slots - 1, n_blocks,
- out->nb[1], out->nb[2], out->nb[3], 0));
- ggml_tensor * score = ggml_scale(ctx0, ggml_mul_mat(ctx0, next, previous), hparams.lilicorr_logit_scale);
- scores = ggml_concat(ctx0, scores,
- ggml_reshape_3d(ctx0, score, top_k * top_k, slots - 1, n_blocks), 1);
- }
-
- ggml_tensor * row = ggml_concat(ctx0, ggml_cast(ctx0, ggml_reshape_3d(ctx0, cand_run, top_k, slots, n_blocks),
- GGML_TYPE_F32), scores, 0);
- ggml_tensor * packed = ggml_pad_ext(ctx0, row, 0, n_embd - row_used, 1, 0, 0, 0, 0, 0);
- packed = ggml_reshape_2d(ctx0, packed, n_embd, n_tokens);
- g.cb(packed, "lilicorr_lattice", -1);
- res->t_h_nextn = packed;
- ggml_build_forward_expand(g.gf, packed);
+ GGML_ASSERT(model_other->output != nullptr && "DFlash draft requires the target model's output projection");
+ return { model_other->output, model_other->output_s };
}
// DFlash decoder, dual-mode by batch type:
// * embd batch -> fused target features: project + inject K/V into the cache.
// * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens
-template <>
-llama_model_dflash::graph<false>::graph(const llm_graph_params & params) : llm_graph_context(params) {
- const int64_t n_embd_inp = hparams.n_embd_inp_enc();
- const int64_t n_embd_head = hparams.n_embd_head_v();
+llama_model_dflash::graph_base::graph_base(const llm_graph_params & params) : llm_graph_context(params) {
+ n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
- ggml_tensor * inp_pos = build_inp_pos();
+ inp_pos = build_inp_pos();
// optional iSWA: pick the matching attention input
- const bool use_iswa = hparams.swa_type != LLAMA_SWA_TYPE_NONE;
-
- llm_graph_input_attn_kv * inp_attn = nullptr;
- llm_graph_input_attn_kv_iswa * inp_attn_iswa = nullptr;
+ use_iswa = hparams.swa_type != LLAMA_SWA_TYPE_NONE;
if (use_iswa) {
inp_attn_iswa = build_attn_inp_kv_iswa();
} else {
inp_attn = build_attn_inp_kv();
}
- const float kq_scale = hparams.f_attention_scale != 0.0f ? hparams.f_attention_scale : 1.0f/sqrtf(float(n_embd_head));
-
// drafts for M-RoPE targets use degenerate sections (temporal dim only)
- int sections[4];
std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);
+}
- auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos) {
- return rope_type == GGML_ROPE_TYPE_MROPE
- ? ggml_rope_multi(ctx0, cur, pos, nullptr,
- n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
- ext_factor, attn_factor, beta_fast, beta_slow)
- : ggml_rope_ext(ctx0, cur, pos, model.layers[0].rope_freqs,
- n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
- ext_factor, attn_factor, beta_fast, beta_slow);
- };
-
- // KV cache injection
- ASSERT_EMBD_OR_TOKEN(ubatch);
- if (ubatch.embd || ubatch.embd_row) {
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
-
- 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));
+ggml_tensor * llama_model_dflash::graph_base::build_rope(ggml_tensor * cur, ggml_tensor * pos) {
+ return rope_type == GGML_ROPE_TYPE_MROPE
+ ? ggml_rope_multi(ctx0, cur, pos, nullptr,
+ n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow)
+ : ggml_rope_ext(ctx0, cur, pos, model.layers[0].rope_freqs,
+ n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+}
- // fuse the target features through the encoder
- ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
- inp_g = build_norm(inp_g, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);
- cb(inp_g, "inp_g_embeddings", -1);
+// KV cache injection
+ggml_tensor * llama_model_dflash::graph_base::build_inject() {
+ ggml_tensor * inp_g = build_inject_inp(*this, model);
+ cb(inp_g, "inp_g_embeddings", -1);
- for (int il = 0; il < n_layer; ++il) {
- const auto & layer = model.layers[il];
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
- ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g, layer.wk_s);
- const bool shared_kv = layer.wv == nullptr;
- ggml_tensor * Vcur = shared_kv ? Kcur : build_lora_mm(layer.wv, inp_g, layer.wv_s);
+ ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g, layer.wk_s);
+ const bool shared_kv = layer.wv == nullptr;
+ ggml_tensor * Vcur = shared_kv ? Kcur : build_lora_mm(layer.wv, inp_g, layer.wv_s);
- Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
- Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
+ Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
+ Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
- Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
- if (shared_kv) {
- Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps);
+ Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
+ if (shared_kv) {
+ Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps);
+ }
+ Kcur = build_rope(Kcur, inp_pos);
+ cb(Kcur, "Kcur_injected", il);
+ cb(Vcur, "Vcur_injected", il);
+
+ if (use_iswa) {
+ // route each layer's K/V to its sub-cache: SWA layers -> sliding cache, full -> dense
+ const bool is_swa = hparams.is_swa(il);
+ const auto * kv = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base();
+ ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs();
+ ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs();
+ // rotate K/V into the cache's rotated space
+ ggml_tensor * k_rot = is_swa ? inp_attn_iswa->self_k_rot_swa[il] : inp_attn_iswa->self_k_rot[il];
+ ggml_tensor * v_rot = is_swa ? inp_attn_iswa->self_v_rot_swa[il] : inp_attn_iswa->self_v_rot[il];
+ if (k_rot) {
+ Kcur = llama_mul_mat_hadamard(ctx0, Kcur, k_rot);
}
- Kcur = build_rope(Kcur, inp_pos);
- cb(Kcur, "Kcur_injected", il);
- cb(Vcur, "Vcur_injected", il);
-
- if (use_iswa) {
- // route each layer's K/V to its sub-cache: SWA layers -> sliding cache, full -> dense
- const bool is_swa = hparams.is_swa(il);
- const auto * kv = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base();
- ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs();
- ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs();
- // rotate K/V into the cache's rotated space
- ggml_tensor * k_rot = is_swa ? inp_attn_iswa->self_k_rot_swa[il] : inp_attn_iswa->self_k_rot[il];
- ggml_tensor * v_rot = is_swa ? inp_attn_iswa->self_v_rot_swa[il] : inp_attn_iswa->self_v_rot[il];
- if (k_rot) {
- Kcur = llama_mul_mat_hadamard(ctx0, Kcur, k_rot);
- }
- if (v_rot) {
- Vcur = llama_mul_mat_hadamard(ctx0, Vcur, v_rot);
- }
- ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il));
- ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il));
- } else {
- // rotate K/V into the cache's rotated space
- if (inp_attn->self_k_rot[il]) {
- Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot[il]);
- }
- if (inp_attn->self_v_rot[il]) {
- Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot[il]);
- }
- ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));
- ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));
+ if (v_rot) {
+ Vcur = llama_mul_mat_hadamard(ctx0, Vcur, v_rot);
}
+ ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il));
+ ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il));
+ } else {
+ // rotate K/V into the cache's rotated space
+ if (inp_attn->self_k_rot[il]) {
+ Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot[il]);
+ }
+ if (inp_attn->self_v_rot[il]) {
+ Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot[il]);
+ }
+ ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));
+ ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));
}
-
- res->t_embd = inp_g;
- if (model.lilicorr_pass_hidden_proj) {
- // Expose context features for retrieval by LiliCorr
- res->t_h_nextn = inp_g;
- }
-
- ggml_build_forward_expand(gf, inp_g);
- return;
}
- // tok_embd from the target model (shared via ctx_other)
- auto * tok_embd = model.tok_embd;
- if (tok_embd == nullptr) {
- GGML_ASSERT(cparams.ctx_other != nullptr);
- const auto * model_other = llama_get_model(cparams.ctx_other);
+ res->t_embd = inp_g;
- GGML_ASSERT(model_other->tok_embd != nullptr && "DFlash decoder requires the target model's token embeddings");
- tok_embd = model_other->tok_embd;
- }
+ ggml_build_forward_expand(gf, inp_g);
+
+ return inp_g;
+}
+ggml_tensor * llama_model_dflash::graph_base::build_inp_tokens() {
ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp_tokens);
res->t_inp_tokens = inp_tokens;
- ggml_tensor * inp_anchor = nullptr;
- if (model.lilicorr_pass_hidden_proj && ubatch.embd_state && n_tokens == (int64_t) hparams.dflash_block_size * ubatch.n_seqs_unq) {
- // LiliCorr sets inp_anchor to pass the context features to the lattice scorer
- auto inp = std::make_unique<llm_graph_input_embd_h>(n_embd, n_embd);
- inp->tokens = inp_tokens;
- inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
- ggml_set_input(inp->h);
- inp_anchor = inp->h;
- res->add_input(std::move(inp));
- } else {
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
- inp->tokens = inp_tokens;
- res->add_input(std::move(inp));
- }
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+ inp->tokens = inp_tokens;
+ res->add_input(std::move(inp));
+
+ return inp_tokens;
+}
+
+ggml_tensor * llama_model_dflash::graph_base::build_layers(ggml_tensor * inp_tokens) {
+ const float kq_scale = hparams.f_attention_scale != 0.0f ? hparams.f_attention_scale : 1.0f/sqrtf(float(n_embd_head));
- ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp_tokens);
+ ggml_tensor * inpL = ggml_get_rows(ctx0, get_tok_embd(*this, model), inp_tokens);
if (hparams.f_embedding_scale != 0.0f) {
inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
}
@@ -1161,17 +453,11 @@ llama_model_dflash::graph<false>::graph(const llm_graph_params & params) : llm_g
res->t_embd = cur;
- // lm_head from the target model (shared via ctx_other)
- auto * output = model.output;
- auto * output_s = model.output_s;
- if (output == nullptr) {
- GGML_ASSERT(cparams.ctx_other != nullptr);
- const auto * model_other = llama_get_model(cparams.ctx_other);
- GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection");
- output = model_other->output;
- output_s = model_other->output_s;
- }
+ return cur;
+}
+ggml_tensor * llama_model_dflash::graph_base::build_logits(ggml_tensor * cur) {
+ const auto [output, output_s] = get_lm_head(*this, model);
cur = build_lora_mm(output, cur, output_s);
// DFlash2 feeds these logits to the selector, so they need the target's output
@@ -1205,628 +491,20 @@ llama_model_dflash::graph<false>::graph(const llm_graph_params & params) : llm_g
cb(cur, "result_output", -1);
res->t_logits = cur;
- // DSpark: bias the draft logits with the Markov head
- if (model.dspark_markov_w1) {
- build_dspark_markov_head(*this, model, inp_tokens);
- }
-
- if (model.dflash_selector_hidden) {
- build_dflash2_selector(*this, model, inp_tokens);
- } else if (model.lilicorr_pass_hidden_proj && ubatch.embd_state && n_tokens == (int64_t) hparams.dflash_block_size * ubatch.n_seqs_unq) {
- build_lilicorr_scores(*this, model, tok_embd, inp_anchor);
- }
-}
-
-// DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above):
-// * embd batch -> project main_x through each stage's wkv and inject K into the ring cache
-// * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads
-llama_model_dflash::graph_dsv4::graph_dsv4(const llm_graph_params & params) :
- llama_model_deepseek4::graph(params, no_build{}) {
- const int64_t n_embd_inp = hparams.n_embd_inp_enc();
- const int64_t n_embd_head = hparams.n_embd_head_k();
- const int64_t n_embd_head_rope = hparams.n_rot();
- const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
-
- ggml_tensor * inp_pos = build_inp_pos();
-
- llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
-
- // KV cache injection: fused target features from the encoder
- ASSERT_EMBD_OR_TOKEN(ubatch);
- if (ubatch.embd || ubatch.embd_row) {
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
-
- 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));
-
- // fuse the target features through the encoder
- ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
- inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
- cb(inp_g, "inp_g_embeddings", -1);
-
- for (int il = 0; il < n_layer; ++il) {
- const auto & layer = model.layers[il];
-
- // main-track KV: kv_norm(wkv(main_x)) with rope on the trailing dims, same
- // rope parameters as the uncompressed layers in build_attention_impl
- ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);
- kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
- kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
-
- kv = ggml_rope_ext(ctx0, kv, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
- freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
- kv = ggml_rope_set_offset(kv, n_embd_head_nope);
- cb(kv, "kv_injected", il);
-
- if (inp_attn->self_k_rot_swa[il]) {
- kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
- }
- ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
- }
-
- res->t_embd = inp_g;
-
- ggml_build_forward_expand(gf, inp_g);
- return;
- }
-
- // tok_embd from the target model (shared via ctx_other)
- auto * tok_embd = model.tok_embd;
- if (tok_embd == nullptr) {
- GGML_ASSERT(cparams.ctx_other != nullptr);
- const auto * model_other = llama_get_model(cparams.ctx_other);
-
- GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");
- tok_embd = model_other->tok_embd;
- }
-
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
-
- inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
- ggml_set_input(inp->tokens);
-
- ggml_tensor * inp_tokens = inp->tokens;
-
- ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
- cb(inpL, "inp_noise_embd", -1);
-
- res->add_input(std::move(inp));
-
- const int64_t hc = hparams.dsv4_hc_mult;
- inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);
- inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
- cb(inpL, "hc_init", -1);
-
- for (int il = 0; il < n_layer; ++il) {
- const auto & layer = model.layers[il];
-
- ggml_tensor * residual = inpL;
- hc_pre attn_hc = build_hc_pre(inpL,
- layer.hc_attn_fn,
- layer.hc_attn_scale,
- layer.hc_attn_base,
- il);
- ggml_tensor * cur = attn_hc.cur;
- cb(cur, "hc_attn_pre", il);
-
- cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
- cb(cur, "attn_norm", il);
-
- cur = build_attention(inp_attn, cur, inp_pos, il);
-
- inpL = build_hc_post(cur, residual, attn_hc.post, attn_hc.comb, il);
- cb(inpL, "hc_attn_post", il);
-
- residual = inpL;
- hc_pre ffn_hc = build_hc_pre(inpL,
- layer.hc_ffn_fn,
- layer.hc_ffn_scale,
- layer.hc_ffn_base,
- il);
- cur = ffn_hc.cur;
- cb(cur, "hc_ffn_pre", il);
-
- cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
- cb(cur, "ffn_norm", il);
-
- ggml_tensor * moe_out = build_moe_ffn(cur,
- layer.ffn_gate_inp,
- layer.ffn_up_exps,
- layer.ffn_gate_exps,
- layer.ffn_down_exps,
- layer.ffn_exp_probs_b,
- n_expert, hparams.n_expert_used(),
- LLM_FFN_SILU, hparams.expert_weights_norm,
- hparams.expert_weights_scale,
- (llama_expert_gating_func_type) hparams.expert_gating_func,
- il);
- cb(moe_out, "ffn_moe_out", il);
-
- ggml_tensor * ffn_shexp = build_ffn(cur,
- layer.ffn_up_shexp, nullptr, nullptr,
- layer.ffn_gate_shexp, nullptr, nullptr,
- layer.ffn_down_shexp, nullptr, nullptr,
- nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
- cb(ffn_shexp, "ffn_shexp", il);
-
- cur = ggml_add(ctx0, moe_out, ffn_shexp);
- cb(cur, "ffn_out", il);
-
- inpL = build_hc_post(cur, residual, ffn_hc.post, ffn_hc.comb, il);
- cb(inpL, "l_out", il);
- }
-
- ggml_tensor * cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
- cb(cur, "hc_head", -1);
-
- // confidence head input: the reference scores the pre-norm collapsed hidden state
- res->t_embd = cur;
-
- cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
- cb(cur, "result_norm", -1);
-
- // lm_head from the target model (shared via ctx_other)
- auto * output = model.output;
- auto * output_s = model.output_s;
- if (output == nullptr) {
- GGML_ASSERT(cparams.ctx_other != nullptr);
- const auto * model_other = llama_get_model(cparams.ctx_other);
- GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
- output = model_other->output;
- output_s = model_other->output_s;
- }
-
- cur = build_lora_mm(output, cur, output_s);
- cb(cur, "result_output", -1);
- res->t_logits = cur;
-
- if (model.dspark_markov_w1) {
- build_dspark_markov_head(*this, model, inp_tokens);
- }
-}
-
-// DeepSeek-V4.1 DSpark decoder (mtp.0..2), dual-mode by batch type like graph_dsv4:
-// * embd batch -> writes every stage's kv_norm(wkv(main_proj(target features))) into its window ring
-// * token batch -> [anchor, noise...] block through the stages, single-pass mHC, no per-head q norm
-// window of a block = window of its first query, see n_swa
-
-llama_model_dflash::graph_dsv41::graph_dsv41(const llm_graph_params & params) :
- llama_model_deepseek41::graph(params, no_build{}) {
- const int64_t n_embd_inp = hparams.n_embd_inp_enc();
- const int64_t n_embd_head = hparams.n_embd_head_k();
- const int64_t n_embd_head_rope = hparams.n_rot();
- const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
-
- GGML_ASSERT(hparams.dsv4_hc_single_pass);
-
- ggml_tensor * inp_pos = build_inp_pos();
-
- llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
-
- // ratio-0 layers: plain rope at rope_theta, no YaRN (original_seq_len 0)
- const auto rope = [&](ggml_tensor * t) {
- t = ggml_rope_ext(ctx0, t, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
- freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
- return ggml_rope_set_offset(t, n_embd_head_nope);
- };
-
- ASSERT_EMBD_OR_TOKEN(ubatch);
- if (ubatch.embd || ubatch.embd_row) {
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
-
- 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));
-
- ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
- inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
- cb(inp_g, "main_x", -1);
-
- for (int il = 0; il < n_layer; ++il) {
- const auto & layer = model.layers[il];
-
- ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);
- kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
- kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
- kv = rope(kv);
- cb(kv, "kv_injected", il);
-
- if (inp_attn->self_k_rot_swa[il]) {
- kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
- }
- ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
- }
-
- res->t_embd = inp_g;
-
- ggml_build_forward_expand(gf, inp_g);
- return;
- }
-
- auto * tok_embd = model.tok_embd;
- if (tok_embd == nullptr) {
- GGML_ASSERT(cparams.ctx_other != nullptr);
- const auto * model_other = llama_get_model(cparams.ctx_other);
- GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");
- tok_embd = model_other->tok_embd;
- }
-
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
-
- inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
- ggml_set_input(inp->tokens);
-
- ggml_tensor * inp_tokens = inp->tokens;
-
- ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
- cb(inpL, "inp_noise_embd", -1);
-
- res->add_input(std::move(inp));
-
- const int64_t hc = hparams.dsv4_hc_mult;
- const int64_t n_groups = hparams.dsv4_o_group_count;
- const int64_t n_heads_group = n_head / n_groups;
- const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
- const int64_t o_group_dim = n_heads_group*n_embd_head;
- const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
-
- GGML_ASSERT(n_embd_head == n_embd_head_v);
- GGML_ASSERT(n_head % n_groups == 0);
-
- inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);
- inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
- cb(inpL, "hc_init", -1);
-
- // the whole block shares the mask row of its first query (the anchor)
- inp_attn->self_attn_mask_swa.block_uniform = true;
-
- ggml_tensor * pre_prev = nullptr;
-
- for (int il = 0; il < n_layer; ++il) {
- const auto & layer = model.layers[il];
-
- ggml_tensor * residual = inpL;
- const hc_mixes attn_mix = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, il);
-
- ggml_tensor * cur;
- if (pre_prev == nullptr) {
- // identity pre-mix of the first stage: copy 0
- cur = ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0));
- } else {
- cur = build_hc_pre(inpL, pre_prev, il);
- }
- cb(cur, "hc_attn_pre", il);
-
- cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
- cb(cur, "attn_norm", il);
-
- // non-causal draft context: block queries see the main-KV ring plus every block key
- {
- ggml_tensor * qr = build_lora_mm(layer.wq_a, cur);
- qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
- cb(qr, "qr_norm", il);
-
- ggml_tensor * q = build_lora_mm(layer.wq_b, qr);
- q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, n_tokens);
- q = rope(q);
- cb(q, "q", il);
-
- ggml_tensor * kv = build_lora_mm(layer.wkv, cur);
- kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
- kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
- kv = rope(kv);
- cb(kv, "kv", il);
-
- ggml_tensor * out = build_attn(inp_attn,
- nullptr, nullptr, nullptr,
- q, kv, kv,
- nullptr, layer.attn_sinks, nullptr,
- kq_scale, il);
- cb(out, "attn_raw", il);
-
- out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, n_tokens);
- out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
- freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
- out = ggml_rope_set_offset(out, n_embd_head_nope);
- cb(out, "attn_derope", il);
-
- out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, n_tokens);
- out = ggml_permute(ctx0, out, 0, 2, 1, 3);
- ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out);
- cb(oa, "attn_wo_a", il);
- oa = ggml_permute(ctx0, oa, 0, 2, 1, 3);
- oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, n_tokens);
-
- cur = build_lora_mm(layer.wo_b, oa);
- cb(cur, "attn_out", il);
- }
-
- inpL = build_hc_post(cur, residual, attn_mix.post, attn_mix.comb, il);
- cb(inpL, "hc_attn_post", il);
-
- residual = inpL;
- const hc_mixes ffn_mix = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, il);
- pre_prev = ffn_mix.pre;
-
- cur = build_hc_pre(inpL, attn_mix.pre, il);
- cb(cur, "hc_ffn_pre", il);
-
- cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
- cb(cur, "ffn_norm", il);
-
- ggml_tensor * moe_out = build_moe_ffn(cur,
- layer.ffn_gate_inp,
- layer.ffn_up_exps,
- layer.ffn_gate_exps,
- layer.ffn_down_exps,
- layer.ffn_exp_probs_b,
- n_expert, hparams.n_expert_used(),
- LLM_FFN_SILU, hparams.expert_weights_norm,
- hparams.expert_weights_scale,
- (llama_expert_gating_func_type) hparams.expert_gating_func,
- il);
- cb(moe_out, "ffn_moe_out", il);
-
- ggml_tensor * ffn_shexp = build_ffn(cur,
- layer.ffn_up_shexp, nullptr, nullptr,
- layer.ffn_gate_shexp, nullptr, nullptr,
- layer.ffn_down_shexp, nullptr, nullptr,
- nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
- cb(ffn_shexp, "ffn_shexp", il);
-
- cur = ggml_add(ctx0, moe_out, ffn_shexp);
- cb(cur, "ffn_out", il);
-
- inpL = build_hc_post(cur, residual, ffn_mix.post, ffn_mix.comb, il);
-
- ggml_tensor * l_out = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
- cb(l_out, "l_out", il);
- ggml_build_forward_expand(gf, l_out);
- }
-
- // no hc_head: collapse with the last FFN's pre-mix
- ggml_tensor * cur = build_hc_pre(inpL, pre_prev, -1);
- cb(cur, "hc_out", -1);
-
- // confidence head takes the un-normed collapsed state
- res->t_embd = cur;
-
- cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
- cb(cur, "result_norm", -1);
-
- auto * output = model.output;
- auto * output_s = model.output_s;
- if (output == nullptr) {
- GGML_ASSERT(cparams.ctx_other != nullptr);
- const auto * model_other = llama_get_model(cparams.ctx_other);
- GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
- output = model_other->output;
- output_s = model_other->output_s;
- }
-
- cur = build_lora_mm(output, cur, output_s);
- cb(cur, "result_output", -1);
- res->t_logits = cur;
-
- if (model.dspark_markov_w1) {
- build_dspark_markov_head(*this, model, inp_tokens);
- }
+ return cur;
}
-// Mistral Large 4 DSpark, dual-mode like graph_dsv41:
-// - embd batch: write latent K of target_norm(fc(features)) to the cache, no attn norm
-// - token batch: decode [anchor, mask_1, ...] block, non-causal
-llama_model_dflash::graph_mistral_large4::graph_mistral_large4(const llm_graph_params & params) : llm_graph_context(params) {
- const int64_t n_embd_inp = hparams.n_embd_inp_enc();
- const int64_t kv_lora_rank = hparams.n_lora_kv;
- const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
- const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
- const int64_t n_embd_head_nope = n_embd_head_k_mla - n_rot;
-
- const float kq_scale = 1.0f/sqrtf(float(n_embd_head_k_mla));
-
- ggml_tensor * inp_pos = build_inp_pos();
-
- llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
-
- // no LeRoPE in the draft
- const auto rope = [&](ggml_tensor * t) {
- return ggml_rope_ext(ctx0, t, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
- ext_factor, attn_factor, beta_fast, beta_slow);
- };
-
- const auto build_kv = [&](ggml_tensor * x, int il) -> std::pair<ggml_tensor *, ggml_tensor *> {
- const auto & layer = model.layers[il];
-
- ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, x);
- cb(kv_cmpr_pe, "kv_cmpr_pe", il);
-
- ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
- ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot), 0);
- kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
- kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
- cb(kv_cmpr, "kv_cmpr", il);
-
- ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_rot, 1, n_tokens,
- ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot),
- ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot),
- ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
- k_pe = rope(k_pe);
- cb(k_pe, "k_pe", il);
-
- return { ggml_concat(ctx0, kv_cmpr, k_pe, 0), kv_cmpr };
- };
-
+llama_model_dflash::graph_dflash::graph_dflash(const llm_graph_params & params) : graph_base(params) {
ASSERT_EMBD_OR_TOKEN(ubatch);
if (ubatch.embd || ubatch.embd_row) {
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
-
- 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));
-
- ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
- inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
- cb(inp_g, "main_x", -1);
-
- for (int il = 0; il < n_layer; ++il) {
- ggml_tensor * kv = build_kv(inp_g, il).first;
- cb(kv, "kv_injected", il);
-
- if (inp_attn->self_k_rot_swa[il]) {
- kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
- }
- ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
- }
-
- res->t_embd = inp_g;
-
- ggml_build_forward_expand(gf, inp_g);
+ build_inject();
return;
}
- auto * tok_embd = model.tok_embd;
- if (tok_embd == nullptr) {
- GGML_ASSERT(cparams.ctx_other != nullptr);
- const auto * model_other = llama_get_model(cparams.ctx_other);
- GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");
- tok_embd = model_other->tok_embd;
- }
-
- auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
-
- inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
- ggml_set_input(inp->tokens);
-
- ggml_tensor * inp_tokens = inp->tokens;
-
- res->add_input(std::move(inp));
-
- const int64_t n_blocks = ubatch.n_seqs_unq;
- const int64_t block_drafts = n_blocks > 0 ? n_tokens / n_blocks : 0;
-
- ggml_tensor * inpL;
- if (block_drafts > 0 && n_tokens % n_blocks == 0 && block_drafts <= (int64_t) hparams.dflash_block_size) {
- // slot k > 0 uses mask row k - 1
- ggml_tensor * anchors = ggml_view_2d(ctx0, inp_tokens, 1, n_blocks, (size_t) block_drafts * inp_tokens->nb[0], 0);
- anchors = ggml_get_rows(ctx0, tok_embd, ggml_cont_1d(ctx0, anchors, n_blocks));
- inpL = ggml_reshape_3d(ctx0, anchors, n_embd, 1, n_blocks);
- if (block_drafts > 1) {
- ggml_tensor * masks = ggml_view_3d(ctx0, model.dspark_mask_embd, n_embd, block_drafts - 1, 1,
- model.dspark_mask_embd->nb[1], model.dspark_mask_embd->nb[2], 0);
- masks = ggml_repeat_4d(ctx0, ggml_cast(ctx0, masks, GGML_TYPE_F32), n_embd, block_drafts - 1, n_blocks, 1);
- inpL = ggml_concat(ctx0, inpL, masks, 1);
- }
- inpL = ggml_reshape_2d(ctx0, inpL, n_embd, n_tokens);
- } else {
- // not a draft block (e.g. graph reserve)
- inpL = ggml_get_rows(ctx0, tok_embd, inp_tokens);
- }
- cb(inpL, "inp_noise_embd", -1);
-
- // the whole block shares the mask row of its first query (the anchor)
- inp_attn->self_attn_mask_swa.block_uniform = true;
-
- for (int il = 0; il < n_layer; ++il) {
- const auto & layer = model.layers[il];
-
- ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
- cb(cur, "attn_norm", il);
-
- {
- ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq, cur);
- cb(q, "q", il);
-
- // {n_embd_head_nope, n_tokens, n_head}
- ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_nope, n_head, n_tokens,
- ggml_row_size(q->type, n_embd_head_k_mla),
- ggml_row_size(q->type, n_embd_head_k_mla) * n_head, 0);
- q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
-
- // {kv_lora_rank, n_head, n_tokens}
- ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
- q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
- cb(q_nope_absorbed, "q_nope_absorbed", il);
-
- ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_rot, n_head, n_tokens,
- ggml_row_size(q->type, n_embd_head_k_mla),
- ggml_row_size(q->type, n_embd_head_k_mla) * n_head,
- ggml_row_size(q->type, n_embd_head_nope));
- q_pe = rope(q_pe);
- cb(q_pe, "q_pe", il);
-
- ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
- cb(Qcur, "Qcur", il);
-
- auto [Kcur, Vcur] = build_kv(cur, il);
- cb(Kcur, "Kcur", il);
-
- // apply wv_b after attn, so the K rotation is already undone
- ggml_tensor * out = build_attn(inp_attn,
- nullptr, nullptr, nullptr,
- Qcur, Kcur, Vcur,
- nullptr, nullptr, nullptr,
- kq_scale, il);
- cb(out, "attn_latent", il);
-
- // {kv_lora_rank, n_tokens, n_head} -> {n_embd_head_v_mla, n_tokens, n_head}
- out = ggml_reshape_3d(ctx0, out, kv_lora_rank, n_head, n_tokens);
- out = ggml_mul_mat(ctx0, layer.wv_b, ggml_permute(ctx0, out, 0, 2, 1, 3));
- out = ggml_cont_2d(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3), n_embd_head_v_mla*n_head, n_tokens);
-
- cur = build_lora_mm(layer.wo, out, layer.wo_s);
- cb(cur, "attn_out", il);
- }
-
- ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
- cb(ffn_inp, "ffn_inp", il);
-
- cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
- cb(cur, "ffn_norm", il);
-
- cur = build_ffn(cur,
- layer.ffn_up, nullptr, nullptr,
- layer.ffn_gate, nullptr, nullptr,
- layer.ffn_down, nullptr, nullptr,
- nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
- cb(cur, "ffn_out", il);
-
- inpL = ggml_add(ctx0, cur, ffn_inp);
- cb(inpL, "l_out", il);
- }
-
- ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1);
- cb(cur, "result_norm", -1);
- res->t_embd = cur;
-
- auto * output = model.output;
- auto * output_s = model.output_s;
- if (output == nullptr) {
- GGML_ASSERT(cparams.ctx_other != nullptr);
- const auto * model_other = llama_get_model(cparams.ctx_other);
- GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
- output = model_other->output;
- output_s = model_other->output_s;
- }
-
- ggml_tensor * logits = build_lora_mm(output, cur, output_s);
-
- // softcap does not apply to the markov bias
- if (hparams.f_final_logit_softcapping) {
- logits = ggml_scale(ctx0, logits, 1.0f / hparams.f_final_logit_softcapping);
- logits = ggml_tanh(ctx0, logits);
- logits = ggml_scale(ctx0, logits, hparams.f_final_logit_softcapping);
- }
- cb(logits, "result_output", -1);
- res->t_logits = logits;
+ ggml_tensor * inp_tokens = build_inp_tokens();
+ ggml_tensor * cur = build_layers(inp_tokens);
- if (model.dspark_markov_w1) {
- build_dspark_markov_head(*this, model, inp_tokens, cur);
- }
+ build_logits(cur);
}
std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const llm_graph_params & params) const {
@@ -1835,16 +513,19 @@ std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const ll
return std::make_unique<graph<true>>(params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
- if (hparams.n_lora_kv > 0) {
- return std::make_unique<graph_mistral_large4>(params);
- }
- if (hparams.dsv4_hc_mult > 0 && hparams.dsv4_hc_single_pass) {
- return std::make_unique<graph_dsv41>(params);
+ switch (backbone) {
+ case BACKBONE_BASE: break;
+ case BACKBONE_DSV4: return std::make_unique<graph_dsv4>(params);
+ case BACKBONE_DSV41: return std::make_unique<graph_dsv41>(params);
+ case BACKBONE_MISTRAL_LARGE4: return std::make_unique<graph_mistral_large4>(params);
}
- if (hparams.dsv4_hc_mult > 0) {
- return std::make_unique<graph_dsv4>(params);
+ switch (variant) {
+ case VARIANT_DFLASH: return std::make_unique<graph_dflash>(params);
+ case VARIANT_DFLASH2: return std::make_unique<graph_dflash2>(params);
+ case VARIANT_DSPARK: return std::make_unique<graph_dspark>(params);
+ case VARIANT_LILICORR: return std::make_unique<graph_lilicorr>(params);
}
- return std::make_unique<graph<false>>(params);
+ GGML_ABORT("invalid DFlash variant");
default:
GGML_ABORT("invalid graph type");
};
diff --git a/src/models/dflash2.cpp b/src/models/dflash2.cpp
new file mode 100644
index 000000000..b384d9527
--- /dev/null
+++ b/src/models/dflash2.cpp
@@ -0,0 +1,224 @@
+#include "models.h"
+
+#include "llama-impl.h"
+
+void llama_model_dflash::load_tensors_dflash2() {
+ LLAMA_LOAD_LOCALS;
+
+ const int64_t rank = hparams.dflash_selector_rank;
+ if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 ||
+ hparams.dflash_conv_kernel_size <= 0 || hparams.dflash_conv_group_size <= 0) {
+ throw std::runtime_error("DFlash2 model is missing conv/selector metadata");
+ }
+ if (n_embd < hparams.dflash_selector_top_k * (hparams.dflash_selector_top_k + 1)) {
+ throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice");
+ }
+
+ dflash_selector_prev = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV, "weight"), { rank, n_vocab }, 0);
+ dflash_selector_next = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT, "weight"), { rank, n_vocab }, 0);
+ dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0);
+
+ load_tensors_dflash2_conv();
+
+ LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__,
+ hparams.dflash_conv_kernel_size, hparams.dflash_conv_group_size,
+ hparams.dflash_selector_rank, hparams.dflash_selector_top_k);
+}
+
+void llama_model_dflash::load_tensors_dflash2_conv() {
+ LLAMA_LOAD_LOCALS;
+
+ const bool conv_meta = hparams.dflash_conv_kernel_size > 0 && hparams.dflash_conv_group_size > 0;
+ if (!conv_meta) {
+ return;
+ }
+ if (n_embd % hparams.dflash_conv_group_size != 0) {
+ throw std::runtime_error("DFlash hidden size must be divisible by conv_group_size");
+ }
+
+ const int64_t kernel = hparams.dflash_conv_kernel_size;
+ const int64_t groups = n_embd / hparams.dflash_conv_group_size;
+ const int64_t projected = 2 * kernel * groups;
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ layer.dflash_attn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);
+ layer.dflash_attn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);
+ layer.dflash_ffn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);
+ layer.dflash_ffn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);
+ }
+}
+
+ggml_tensor * llama_model_dflash::build_dflash2_conv(
+ llm_graph_context & g,
+ ggml_tensor * hidden,
+ ggml_tensor * dynamic,
+ ggml_tensor * base,
+ int side) {
+ const auto & hparams = g.hparams;
+ const int64_t hidden_size = hidden->ne[0];
+ const int64_t n_tokens = hidden->ne[1];
+ const int64_t n_blocks = g.ubatch.n_seqs_unq;
+ const int64_t kernel_size = hparams.dflash_conv_kernel_size;
+ const int64_t group_size = hparams.dflash_conv_group_size;
+ const int64_t n_groups = hidden_size / group_size;
+
+ GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
+ GGML_ASSERT(dynamic && base && side >= 0 && side < 2);
+
+ const int64_t block_size = n_tokens / n_blocks;
+ ggml_context * ctx0 = g.ctx0;
+ // ggml_cont copies even when the tensor is already contiguous
+ if (!ggml_is_contiguous(hidden) || hidden->ne[1] != n_tokens) {
+ hidden = ggml_cont_2d(ctx0, hidden, hidden_size, n_tokens);
+ }
+ if (!ggml_is_contiguous(dynamic) || dynamic->ne[1] != n_tokens) {
+ dynamic = ggml_cont_2d(ctx0, dynamic, dynamic->ne[0], n_tokens);
+ }
+ ggml_tensor * blocks = ggml_reshape_3d(ctx0, hidden, hidden_size, block_size, n_blocks);
+ ggml_tensor * coeffs = ggml_reshape_4d(ctx0, dynamic, n_groups, kernel_size, 2, n_tokens);
+ ggml_tensor * coeffs_side = ggml_view_3d(ctx0, coeffs, n_groups, kernel_size, n_tokens,
+ coeffs->nb[1], coeffs->nb[3], side * coeffs->nb[2]);
+
+ ggml_tensor * coeff_all = ggml_cont(ctx0, coeffs_side);
+ coeff_all = ggml_reshape_4d(ctx0, coeff_all, 1, n_groups, kernel_size, n_tokens);
+ coeff_all = ggml_repeat_4d(ctx0, coeff_all, group_size, n_groups, kernel_size, n_tokens);
+
+ ggml_tensor * base_side = ggml_reshape_4d(ctx0,
+ ggml_view_1d(ctx0, base, hidden_size * kernel_size, side * base->nb[2]),
+ group_size, n_groups, kernel_size, 1);
+
+ ggml_tensor * weight_all = ggml_add(ctx0, coeff_all, base_side);
+
+ // taps at or past block_size only read the left padding and add nothing
+ const int64_t n_taps = std::min(kernel_size, block_size);
+
+ ggml_tensor * result = nullptr;
+ for (int64_t tap = 0; tap < n_taps; ++tap) {
+ ggml_tensor * values = blocks;
+ if (tap > 0) {
+ ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks,
+ blocks->nb[1], blocks->nb[2], 0);
+ values = ggml_pad_ext(ctx0, previous, 0, 0, tap, 0, 0, 0, 0, 0);
+ }
+ values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens);
+
+ ggml_tensor * weight = ggml_reshape_2d(ctx0,
+ ggml_cont(ctx0, ggml_view_4d(ctx0, weight_all, group_size, n_groups, 1, n_tokens,
+ weight_all->nb[1], weight_all->nb[2], weight_all->nb[3], tap * weight_all->nb[2])),
+ hidden_size, n_tokens);
+
+ ggml_tensor * term = ggml_mul(ctx0, weight, values);
+ result = result ? ggml_add(ctx0, result, term) : term;
+ }
+ return result;
+}
+
+// DFlash2 selector: top-k candidates per block position plus the pairwise
+// transition scores, packed into the nextn output slot for the CPU-side walk.
+static void build_dflash2_selector(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {
+ ggml_context * ctx0 = g.ctx0;
+ auto & res = g.res;
+
+ const auto & hparams = g.hparams;
+ const int64_t n_tokens = g.n_tokens;
+ const int64_t n_embd = g.n_embd;
+
+ const int64_t top_k = hparams.dflash_selector_top_k;
+ const int64_t rank = hparams.dflash_selector_rank;
+ const int64_t n_blocks = g.ubatch.n_seqs_unq;
+ GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
+ GGML_ASSERT(res->t_logits->ne[1] == n_tokens);
+ if (!tokens) {
+ return;
+ }
+
+ const int64_t tokens_per_block = n_tokens / n_blocks;
+ const int64_t block_size = std::min<int64_t>(tokens_per_block, hparams.dflash_block_size);
+ const int64_t row_used = top_k + top_k * top_k;
+
+ ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k);
+ ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens);
+ ggml_tensor * unary = ggml_reshape_2d(ctx0,
+ ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens);
+ ggml_tensor * gate = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd);
+
+ // Everything below indexes [.., tokens_per_block, n_blocks]: the block
+ // position varies fastest, sequences are the outer dimension.
+ ggml_tensor * cand_blk = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks);
+ ggml_tensor * unary_blk = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks);
+ ggml_tensor * gate_blk = ggml_reshape_3d(ctx0, gate, rank, tokens_per_block, n_blocks);
+
+ // a position's score reads only the candidate sets at pos-1 and pos, so a run
+ // of positions has no internal dependency and scores in one batched matmul
+ auto score_run = [&](int64_t beg_pos, int64_t n_pos, ggml_tensor * pred_ids) {
+ ggml_tensor * cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, n_pos, n_blocks,
+ cand_blk->nb[1], cand_blk->nb[2], beg_pos * cand_blk->nb[1]));
+ ggml_tensor * unary_run = ggml_cont(ctx0, ggml_view_3d(ctx0, unary_blk, top_k, n_pos, n_blocks,
+ unary_blk->nb[1], unary_blk->nb[2], beg_pos * unary_blk->nb[1]));
+ ggml_tensor * gate_run = ggml_cont(ctx0, ggml_view_3d(ctx0, gate_blk, rank, n_pos, n_blocks,
+ gate_blk->nb[1], gate_blk->nb[2], beg_pos * gate_blk->nb[1]));
+
+ const int64_t n_pred = pred_ids->ne[0] / (n_pos * n_blocks);
+
+ ggml_tensor * successor = ggml_reshape_4d(ctx0,
+ ggml_get_rows(ctx0, model.dflash_selector_next, ggml_reshape_1d(ctx0, cand_run, top_k * n_pos * n_blocks)),
+ rank, top_k, n_pos, n_blocks);
+ ggml_tensor * predecessor = ggml_reshape_4d(ctx0,
+ ggml_get_rows(ctx0, model.dflash_selector_prev, pred_ids),
+ rank, n_pred, n_pos, n_blocks);
+
+ ggml_tensor * gate_bcast = ggml_reshape_4d(ctx0, gate_run, rank, 1, n_pos, n_blocks);
+ ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, gate_bcast, predecessor));
+ ggml_tensor * score = ggml_mul_mat(ctx0, successor, cond);
+ if (n_pred == 1) {
+ score = ggml_repeat_4d(ctx0, score, top_k, top_k, n_pos, n_blocks);
+ }
+ ggml_tensor * unary_bcast = ggml_reshape_4d(ctx0, unary_run, top_k, 1, n_pos, n_blocks);
+ score = ggml_add(ctx0, score, ggml_repeat(ctx0, unary_bcast, score));
+
+ ggml_tensor * row = ggml_concat(ctx0,
+ ggml_cast(ctx0, cand_run, GGML_TYPE_F32),
+ ggml_reshape_3d(ctx0, score, top_k * top_k, n_pos, n_blocks), 0);
+ return ggml_pad(ctx0, row, n_embd - row_used, 0, 0, 0);
+ };
+
+ ggml_tensor * packed = ggml_fill(ctx0,
+ ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f);
+
+ if (block_size > 1) {
+ // Position 1 alone: its predecessor is the anchor token, one id per
+ // sequence rather than a candidate set.
+ ggml_tensor * anchor_ids = ggml_cont_1d(ctx0,
+ ggml_view_2d(ctx0, tokens, 1, n_blocks, tokens_per_block * tokens->nb[0], 0), n_blocks);
+ packed = ggml_concat(ctx0, packed, score_run(1, 1, anchor_ids), 1);
+ }
+ if (block_size > 2) {
+ ggml_tensor * prev_ids = ggml_reshape_1d(ctx0,
+ ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, block_size - 2, n_blocks,
+ cand_blk->nb[1], cand_blk->nb[2], cand_blk->nb[1])),
+ top_k * (block_size - 2) * n_blocks);
+ packed = ggml_concat(ctx0, packed, score_run(2, block_size - 2, prev_ids), 1);
+ }
+
+ packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks);
+ g.cb(packed, "dflash2_lattice", -1);
+ res->t_h_nextn = packed;
+ ggml_build_forward_expand(g.gf, packed);
+}
+
+llama_model_dflash::graph_dflash2::graph_dflash2(const llm_graph_params & params) : graph_base(params) {
+ ASSERT_EMBD_OR_TOKEN(ubatch);
+ if (ubatch.embd || ubatch.embd_row) {
+ build_inject();
+ return;
+ }
+
+ ggml_tensor * inp_tokens = build_inp_tokens();
+ ggml_tensor * cur = build_layers(inp_tokens);
+
+ build_logits(cur);
+
+ build_dflash2_selector(*this, model, inp_tokens);
+}
diff --git a/src/models/dspark.cpp b/src/models/dspark.cpp
new file mode 100644
index 000000000..b92c89761
--- /dev/null
+++ b/src/models/dspark.cpp
@@ -0,0 +1,873 @@
+#include "models.h"
+
+#include "llama-impl.h"
+#include "llama-kv-cache.h"
+#include "llama-kv-cache-iswa.h"
+
+// DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
+void llama_model_dflash::load_hparams_dsv4(llama_model_loader & ml) {
+ ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
+ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
+ ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
+ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm);
+ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
+ ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all);
+ if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {
+ hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;
+ }
+ ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count);
+ ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank);
+ ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
+ ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
+ ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false);
+ ml.get_key(LLM_KV_HYPER_CONNECTION_SINGLE_PASS, hparams.dsv4_hc_single_pass, false);
+
+ GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero
+
+ if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
+ throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring");
+ }
+ for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
+ if (hparams.dsv4_compress_ratios[il] != 0) {
+ throw std::runtime_error("DSpark DSV4 draft expects uncompressed attention on all stages");
+ }
+ }
+
+ GGML_ASSERT(hparams.n_swa > 0);
+ if (hparams.dsv4_hc_single_pass) {
+ // DeepSeek-V4.1 DSpark: block queries see the main-KV ring plus the block via the mask row of the
+ // first query, so that row's SWA window must reach window_size positions before the anchor
+ hparams.n_swa += 1;
+ }
+ hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+ hparams.set_swa_pattern(0);
+ for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
+ hparams.is_swa_impl[il] = true;
+ }
+ hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
+ hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
+}
+
+// Mistral Large 4 DSpark: dense MLA stages, one sliding window, plain YaRN
+void llama_model_dflash::load_hparams_mistral_large4(llama_model_loader & ml) {
+ ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
+ ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
+ ml.get_key(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul, false);
+ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
+
+ if (hparams.n_swa == 0) {
+ throw std::runtime_error("DSpark Mistral Large 4 draft expects a sliding window");
+ }
+ if (hparams.n_embd_head_k_full != hparams.n_lora_kv + hparams.n_rot()) {
+ throw std::runtime_error("DSpark Mistral Large 4 draft expects key_length == kv_lora_rank + rope dim");
+ }
+
+ // same as DSV4.1 DSpark: block uses the window of the anchor
+ hparams.n_swa += 1;
+ hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+ hparams.set_swa_pattern(0);
+ for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
+ hparams.is_swa_impl[il] = true;
+ }
+ hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
+ hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
+}
+
+// DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
+void llama_model_dflash::load_tensors_dspark(int64_t n_vocab_draft) {
+ LLAMA_LOAD_LOCALS;
+
+ const int64_t dspark_markov_rank = ml->get_tensor_meta("markov_w1.weight")->ne[0];
+
+ dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);
+ dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0);
+ dspark_markov_w2_s = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
+
+ dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED);
+ dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED);
+
+ dspark_joint_proj = create_tensor(tn(LLM_TENSOR_DSPARK_JOINT_PROJ, "weight"), { n_embd + 2*dspark_markov_rank, 3*dspark_markov_rank }, TENSOR_NOT_REQUIRED);
+ dspark_joint_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_JOINT_PROJ, "bias"), { 3*dspark_markov_rank }, dspark_joint_proj ? 0 : TENSOR_NOT_REQUIRED);
+
+ LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);
+}
+
+void llama_model_dflash::load_tensors_layers_dsv4() {
+ LLAMA_LOAD_LOCALS;
+
+ const int64_t q_lora_rank = hparams.n_lora_q;
+ const int64_t n_ff_exp = hparams.n_ff_exp();
+ const int64_t n_expert_shared = hparams.n_expert_shared;
+ const int64_t n_embd_head = hparams.n_embd_head_k();
+ const int64_t o_groups = hparams.dsv4_o_group_count;
+ const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
+ const int64_t hc_mult = hparams.dsv4_hc_mult;
+ const int64_t hc_dim = hc_mult * n_embd;
+ const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
+
+ // single-pass mHC: collapse with the last FFN pre-mix, no hc_head
+ if (!hparams.dsv4_hc_single_pass) {
+ hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0);
+ hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0);
+ hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);
+ }
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+ layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);
+ layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
+ layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);
+ layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);
+ layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);
+ layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);
+ layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, TENSOR_ALLOW_RESHAPE);
+ layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
+
+ layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
+ layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);
+ layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
+ layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
+ layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);
+ layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);
+
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
+
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+ }
+}
+
+void llama_model_dflash::load_tensors_layers_mistral_large4() {
+ LLAMA_LOAD_LOCALS;
+
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
+ const int64_t n_embd_head_nope = n_embd_head_k_mla - n_rot;
+
+ dspark_mask_embd = create_tensor(tn(LLM_TENSOR_DSPARK_MASK_EMBD, "weight"), { n_embd, hparams.dflash_block_size - 1 }, 0);
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
+ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_head * n_embd_head_k_mla }, 0);
+ layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), { n_embd, kv_lora_rank + n_rot }, 0);
+ layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), { kv_lora_rank }, 0);
+ layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), { n_embd_head_nope, kv_lora_rank, n_head }, 0);
+ layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), { kv_lora_rank, n_embd_head_v_mla, n_head }, 0);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_head * n_embd_head_v_mla, n_embd }, 0);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0);
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0);
+ }
+}
+
+// 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;
+ auto & res = g.res;
+
+ ggml_tensor * w1 = model.dspark_markov_w1;
+ ggml_tensor * w2 = model.dspark_markov_w2;
+ GGML_ASSERT(w1 && w2 && "DSpark markov weights not loaded");
+
+ // confidence head is optional
+ const bool has_conf = model.dspark_conf_proj != nullptr;
+
+ // rnn head: gated state carried across block positions
+ const bool is_rnn = model.dspark_joint_proj != nullptr;
+ GGML_ASSERT(!is_rnn || hidden != nullptr);
+
+ ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens]
+ const int64_t n_vocab = base->ne[0];
+ const int64_t n_tok = base->ne[1];
+
+ const auto it = model.gguf_kv.find("dflash.block_size");
+ GGML_ASSERT(it != model.gguf_kv.end() && "DSpark draft requires 'dflash.block_size' in GGUF metadata");
+ const int64_t block_size = std::stoi(it->second);
+ GGML_ASSERT(block_size > 0);
+
+ // bonus anchor (SpecForge exports): slot 0 is a bonus token, not a prediction slot
+ const auto it_anchor = model.gguf_kv.find("dflash.sample_from_anchor");
+ const bool sample_from_anchor = it_anchor == model.gguf_kv.end() || it_anchor->second == "true";
+ const int64_t i_draft_beg = sample_from_anchor ? 0 : 1;
+
+ const int64_t n_blocks = g.ubatch.n_seqs_unq;
+ GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks");
+ // runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size
+ const int64_t block_drafts = n_tok / n_blocks;
+ if (block_drafts > block_size) {
+ return;
+ }
+
+ // anchor (committed last) token of every block: token 0 of each block, i.e. a strided view
+ const size_t token_stride = (size_t) block_drafts * tokens->nb[0];
+ const size_t base_stride = (size_t) block_drafts * base->nb[1];
+
+ ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0);
+ prev = ggml_cont_1d(ctx0, prev, n_blocks);
+
+ ggml_tensor * cat = nullptr;
+ ggml_tensor * cat_conf = nullptr;
+
+ if (!sample_from_anchor) {
+ // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column
+ cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));
+ if (has_conf) {
+ cat_conf = ggml_sigmoid(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0));
+ }
+ }
+
+ ggml_tensor * state = nullptr; // [R, n_blocks]
+
+ // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final
+ // token pick, not the Markov conditioning path
+ for (int64_t i = i_draft_beg; i < block_drafts; ++i) {
+ ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks]
+ ggml_tensor * markov = w1_prev;
+
+ if (is_rnn) {
+ const int64_t rank = w1_prev->ne[0];
+ if (state == nullptr) {
+ state = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, rank, n_blocks), 0.0f);
+ }
+
+ ggml_tensor * hidden_i = ggml_view_2d(ctx0, hidden, hidden->ne[0], n_blocks,
+ (size_t) block_drafts * hidden->nb[1], i*hidden->nb[1]);
+ ggml_tensor * joint = ggml_concat(ctx0, ggml_concat(ctx0, state, w1_prev, 0), ggml_cont(ctx0, hidden_i), 0);
+ joint = ggml_add(ctx0, ggml_mul_mat(ctx0, model.dspark_joint_proj, joint), model.dspark_joint_proj_b); // [3R, n_blocks]
+
+ const auto chunk = [&](int64_t k) {
+ return ggml_cont(ctx0, ggml_view_2d(ctx0, joint, rank, n_blocks, joint->nb[1], k*rank*ggml_element_size(joint)));
+ };
+ ggml_tensor * gate = ggml_sigmoid(ctx0, chunk(0));
+ ggml_tensor * cand = ggml_tanh(ctx0, chunk(1));
+
+ // state = gate * state + (1 - gate) * cand
+ state = ggml_add(ctx0, cand, ggml_mul(ctx0, gate, ggml_sub(ctx0, state, cand)));
+ markov = ggml_tanh(ctx0, chunk(2));
+ }
+
+ ggml_tensor * bias = g.build_lora_mm(w2, markov, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks]
+ if (model.d2t) {
+ // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others)
+ const int64_t n_draft_vocab = bias->ne[0];
+ ggml_tensor * full = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_blocks), 0.0f);
+ bias = ggml_set_rows(ctx0, full,
+ ggml_reshape_3d(ctx0, bias, 1, n_draft_vocab, n_blocks),
+ ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1));
+ bias = ggml_reshape_2d(ctx0, bias, n_vocab, n_blocks);
+ }
+
+ // position i of every block: strided view [n_vocab, n_blocks]
+ ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]);
+ ggml_tensor * col = ggml_add(ctx0, base_i, bias);
+
+ cat = cat ? ggml_concat(ctx0, cat, col, 1) : col;
+
+ if (has_conf) {
+ // confidence head input: predicts per-position acceptance
+ ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok]
+ // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks]
+ ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,
+ (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);
+ ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), markov, 0);
+ ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);
+ if (model.dspark_conf_proj_b) {
+ conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);
+ }
+ conf = ggml_sigmoid(ctx0, conf);
+
+ cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;
+ }
+
+ if (i + 1 < block_drafts) {
+ prev = ggml_argmax(ctx0, col);
+ }
+ }
+
+ // cat is position-major; restore ubatch block-major order
+ ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, block_drafts);
+ out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks]
+ out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok);
+
+ if (has_conf) {
+ ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts);
+ conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3));
+ conf = ggml_reshape_2d(ctx0, conf, 1, n_tok);
+
+ // note: broadcast the [1, n_tok] confidences to n_embd-wide rows to be able to reuse `llama_get_embeddings_nextn`
+ conf = ggml_repeat(ctx0, conf, res->t_embd);
+ res->t_h_nextn = conf;
+ ggml_build_forward_expand(g.gf, conf);
+ }
+
+ res->t_logits = out;
+ ggml_build_forward_expand(g.gf, out);
+}
+
+llama_model_dflash::graph_dspark::graph_dspark(const llm_graph_params & params) : graph_base(params) {
+ ASSERT_EMBD_OR_TOKEN(ubatch);
+ if (ubatch.embd || ubatch.embd_row) {
+ build_inject();
+ return;
+ }
+
+ ggml_tensor * inp_tokens = build_inp_tokens();
+ ggml_tensor * cur = build_layers(inp_tokens);
+
+ build_logits(cur);
+
+ // bias the draft logits with the Markov head
+ build_dspark_markov_head(*this, model, inp_tokens);
+}
+
+// DSV4 DSpark decoder, dual-mode by batch type (see graph_base):
+// * embd batch -> project main_x through each stage's wkv and inject K into the ring cache
+// * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads
+llama_model_dflash::graph_dsv4::graph_dsv4(const llm_graph_params & params) :
+ llama_model_deepseek4::graph(params, no_build{}) {
+ const int64_t n_embd_head = hparams.n_embd_head_k();
+ const int64_t n_embd_head_rope = hparams.n_rot();
+ const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+ // KV cache injection: fused target features from the encoder
+ ASSERT_EMBD_OR_TOKEN(ubatch);
+ if (ubatch.embd || ubatch.embd_row) {
+ ggml_tensor * inp_g = build_inject_inp(*this, model);
+ cb(inp_g, "inp_g_embeddings", -1);
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ // main-track KV: kv_norm(wkv(main_x)) with rope on the trailing dims, same
+ // rope parameters as the uncompressed layers in build_attention_impl
+ ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);
+ kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
+ kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
+
+ kv = ggml_rope_ext(ctx0, kv, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
+ freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+ kv = ggml_rope_set_offset(kv, n_embd_head_nope);
+ cb(kv, "kv_injected", il);
+
+ if (inp_attn->self_k_rot_swa[il]) {
+ kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
+ }
+ ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
+ }
+
+ res->t_embd = inp_g;
+
+ ggml_build_forward_expand(gf, inp_g);
+ return;
+ }
+
+ ggml_tensor * tok_embd = get_tok_embd(*this, model);
+
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp->tokens);
+
+ ggml_tensor * inp_tokens = inp->tokens;
+
+ ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
+ cb(inpL, "inp_noise_embd", -1);
+
+ res->add_input(std::move(inp));
+
+ const int64_t hc = hparams.dsv4_hc_mult;
+ inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);
+ inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
+ cb(inpL, "hc_init", -1);
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * residual = inpL;
+ hc_pre attn_hc = build_hc_pre(inpL,
+ layer.hc_attn_fn,
+ layer.hc_attn_scale,
+ layer.hc_attn_base,
+ il);
+ ggml_tensor * cur = attn_hc.cur;
+ cb(cur, "hc_attn_pre", il);
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ cur = build_attention(inp_attn, cur, inp_pos, il);
+
+ inpL = build_hc_post(cur, residual, attn_hc.post, attn_hc.comb, il);
+ cb(inpL, "hc_attn_post", il);
+
+ residual = inpL;
+ hc_pre ffn_hc = build_hc_pre(inpL,
+ layer.hc_ffn_fn,
+ layer.hc_ffn_scale,
+ layer.hc_ffn_base,
+ il);
+ cur = ffn_hc.cur;
+ cb(cur, "hc_ffn_pre", il);
+
+ cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ ggml_tensor * moe_out = build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ layer.ffn_gate_exps,
+ layer.ffn_down_exps,
+ layer.ffn_exp_probs_b,
+ n_expert, hparams.n_expert_used(),
+ LLM_FFN_SILU, hparams.expert_weights_norm,
+ hparams.expert_weights_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il);
+ cb(moe_out, "ffn_moe_out", il);
+
+ ggml_tensor * ffn_shexp = build_ffn(cur,
+ layer.ffn_up_shexp, nullptr, nullptr,
+ layer.ffn_gate_shexp, nullptr, nullptr,
+ layer.ffn_down_shexp, nullptr, nullptr,
+ nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "ffn_out", il);
+
+ inpL = build_hc_post(cur, residual, ffn_hc.post, ffn_hc.comb, il);
+ cb(inpL, "l_out", il);
+ }
+
+ ggml_tensor * cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
+ cb(cur, "hc_head", -1);
+
+ // confidence head input: the reference scores the pre-norm collapsed hidden state
+ res->t_embd = cur;
+
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+
+ const auto [output, output_s] = get_lm_head(*this, model);
+ cur = build_lora_mm(output, cur, output_s);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ if (model.dspark_markov_w1) {
+ build_dspark_markov_head(*this, model, inp_tokens);
+ }
+}
+
+// DeepSeek-V4.1 DSpark decoder (mtp.0..2), dual-mode by batch type like graph_dsv4:
+// * embd batch -> writes every stage's kv_norm(wkv(main_proj(target features))) into its window ring
+// * token batch -> [anchor, noise...] block through the stages, single-pass mHC, no per-head q norm
+// window of a block = window of its first query, see n_swa
+
+llama_model_dflash::graph_dsv41::graph_dsv41(const llm_graph_params & params) :
+ llama_model_deepseek41::graph(params, no_build{}) {
+ const int64_t n_embd_head = hparams.n_embd_head_k();
+ const int64_t n_embd_head_rope = hparams.n_rot();
+ const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
+
+ GGML_ASSERT(hparams.dsv4_hc_single_pass);
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+ // ratio-0 layers: plain rope at rope_theta, no YaRN (original_seq_len 0)
+ const auto rope = [&](ggml_tensor * t) {
+ t = ggml_rope_ext(ctx0, t, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
+ freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+ return ggml_rope_set_offset(t, n_embd_head_nope);
+ };
+
+ ASSERT_EMBD_OR_TOKEN(ubatch);
+ if (ubatch.embd || ubatch.embd_row) {
+ ggml_tensor * inp_g = build_inject_inp(*this, model);
+ cb(inp_g, "main_x", -1);
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);
+ kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
+ kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
+ kv = rope(kv);
+ cb(kv, "kv_injected", il);
+
+ if (inp_attn->self_k_rot_swa[il]) {
+ kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
+ }
+ ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
+ }
+
+ res->t_embd = inp_g;
+
+ ggml_build_forward_expand(gf, inp_g);
+ return;
+ }
+
+ ggml_tensor * tok_embd = get_tok_embd(*this, model);
+
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp->tokens);
+
+ ggml_tensor * inp_tokens = inp->tokens;
+
+ ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
+ cb(inpL, "inp_noise_embd", -1);
+
+ res->add_input(std::move(inp));
+
+ const int64_t hc = hparams.dsv4_hc_mult;
+ const int64_t n_groups = hparams.dsv4_o_group_count;
+ const int64_t n_heads_group = n_head / n_groups;
+ const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
+ const int64_t o_group_dim = n_heads_group*n_embd_head;
+ const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
+
+ GGML_ASSERT(n_embd_head == n_embd_head_v);
+ GGML_ASSERT(n_head % n_groups == 0);
+
+ inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);
+ inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
+ cb(inpL, "hc_init", -1);
+
+ // the whole block shares the mask row of its first query (the anchor)
+ inp_attn->self_attn_mask_swa.block_uniform = true;
+
+ ggml_tensor * pre_prev = nullptr;
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * residual = inpL;
+ const hc_mixes attn_mix = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, il);
+
+ ggml_tensor * cur;
+ if (pre_prev == nullptr) {
+ // identity pre-mix of the first stage: copy 0
+ cur = ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0));
+ } else {
+ cur = build_hc_pre(inpL, pre_prev, il);
+ }
+ cb(cur, "hc_attn_pre", il);
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ // non-causal draft context: block queries see the main-KV ring plus every block key
+ {
+ ggml_tensor * qr = build_lora_mm(layer.wq_a, cur);
+ qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(qr, "qr_norm", il);
+
+ ggml_tensor * q = build_lora_mm(layer.wq_b, qr);
+ q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, n_tokens);
+ q = rope(q);
+ cb(q, "q", il);
+
+ ggml_tensor * kv = build_lora_mm(layer.wkv, cur);
+ kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
+ kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
+ kv = rope(kv);
+ cb(kv, "kv", il);
+
+ ggml_tensor * out = build_attn(inp_attn,
+ nullptr, nullptr, nullptr,
+ q, kv, kv,
+ nullptr, layer.attn_sinks, nullptr,
+ kq_scale, il);
+ cb(out, "attn_raw", il);
+
+ out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, n_tokens);
+ out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
+ freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+ out = ggml_rope_set_offset(out, n_embd_head_nope);
+ cb(out, "attn_derope", il);
+
+ out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, n_tokens);
+ out = ggml_permute(ctx0, out, 0, 2, 1, 3);
+ ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out);
+ cb(oa, "attn_wo_a", il);
+ oa = ggml_permute(ctx0, oa, 0, 2, 1, 3);
+ oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, n_tokens);
+
+ cur = build_lora_mm(layer.wo_b, oa);
+ cb(cur, "attn_out", il);
+ }
+
+ inpL = build_hc_post(cur, residual, attn_mix.post, attn_mix.comb, il);
+ cb(inpL, "hc_attn_post", il);
+
+ residual = inpL;
+ const hc_mixes ffn_mix = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, il);
+ pre_prev = ffn_mix.pre;
+
+ cur = build_hc_pre(inpL, attn_mix.pre, il);
+ cb(cur, "hc_ffn_pre", il);
+
+ cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ ggml_tensor * moe_out = build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ layer.ffn_gate_exps,
+ layer.ffn_down_exps,
+ layer.ffn_exp_probs_b,
+ n_expert, hparams.n_expert_used(),
+ LLM_FFN_SILU, hparams.expert_weights_norm,
+ hparams.expert_weights_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il);
+ cb(moe_out, "ffn_moe_out", il);
+
+ ggml_tensor * ffn_shexp = build_ffn(cur,
+ layer.ffn_up_shexp, nullptr, nullptr,
+ layer.ffn_gate_shexp, nullptr, nullptr,
+ layer.ffn_down_shexp, nullptr, nullptr,
+ nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "ffn_out", il);
+
+ inpL = build_hc_post(cur, residual, ffn_mix.post, ffn_mix.comb, il);
+
+ ggml_tensor * l_out = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+ cb(l_out, "l_out", il);
+ ggml_build_forward_expand(gf, l_out);
+ }
+
+ // no hc_head: collapse with the last FFN's pre-mix
+ ggml_tensor * cur = build_hc_pre(inpL, pre_prev, -1);
+ cb(cur, "hc_out", -1);
+
+ // confidence head takes the un-normed collapsed state
+ res->t_embd = cur;
+
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+
+ const auto [output, output_s] = get_lm_head(*this, model);
+ cur = build_lora_mm(output, cur, output_s);
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+
+ if (model.dspark_markov_w1) {
+ build_dspark_markov_head(*this, model, inp_tokens);
+ }
+}
+
+// Mistral Large 4 DSpark, dual-mode like graph_dsv41:
+// - embd batch: write latent K of target_norm(fc(features)) to the cache, no attn norm
+// - token batch: decode [anchor, mask_1, ...] block, non-causal
+llama_model_dflash::graph_mistral_large4::graph_mistral_large4(const llm_graph_params & params) : llm_graph_context(params) {
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
+ const int64_t n_embd_head_nope = n_embd_head_k_mla - n_rot;
+
+ const float kq_scale = 1.0f/sqrtf(float(n_embd_head_k_mla));
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+ // no LeRoPE in the draft
+ const auto rope = [&](ggml_tensor * t) {
+ return ggml_rope_ext(ctx0, t, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ };
+
+ const auto build_kv = [&](ggml_tensor * x, int il) -> std::pair<ggml_tensor *, ggml_tensor *> {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, x);
+ cb(kv_cmpr_pe, "kv_cmpr_pe", il);
+
+ ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot), 0);
+ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+ cb(kv_cmpr, "kv_cmpr", il);
+
+ ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_rot, 1, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
+ k_pe = rope(k_pe);
+ cb(k_pe, "k_pe", il);
+
+ return { ggml_concat(ctx0, kv_cmpr, k_pe, 0), kv_cmpr };
+ };
+
+ ASSERT_EMBD_OR_TOKEN(ubatch);
+ if (ubatch.embd || ubatch.embd_row) {
+ ggml_tensor * inp_g = build_inject_inp(*this, model);
+ cb(inp_g, "main_x", -1);
+
+ for (int il = 0; il < n_layer; ++il) {
+ ggml_tensor * kv = build_kv(inp_g, il).first;
+ cb(kv, "kv_injected", il);
+
+ if (inp_attn->self_k_rot_swa[il]) {
+ kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa[il]);
+ }
+ ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
+ }
+
+ res->t_embd = inp_g;
+
+ ggml_build_forward_expand(gf, inp_g);
+ return;
+ }
+
+ ggml_tensor * tok_embd = get_tok_embd(*this, model);
+
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp->tokens);
+
+ ggml_tensor * inp_tokens = inp->tokens;
+
+ res->add_input(std::move(inp));
+
+ const int64_t n_blocks = ubatch.n_seqs_unq;
+ const int64_t block_drafts = n_blocks > 0 ? n_tokens / n_blocks : 0;
+
+ ggml_tensor * inpL;
+ if (block_drafts > 0 && n_tokens % n_blocks == 0 && block_drafts <= (int64_t) hparams.dflash_block_size) {
+ // slot k > 0 uses mask row k - 1
+ ggml_tensor * anchors = ggml_view_2d(ctx0, inp_tokens, 1, n_blocks, (size_t) block_drafts * inp_tokens->nb[0], 0);
+ anchors = ggml_get_rows(ctx0, tok_embd, ggml_cont_1d(ctx0, anchors, n_blocks));
+ inpL = ggml_reshape_3d(ctx0, anchors, n_embd, 1, n_blocks);
+ if (block_drafts > 1) {
+ ggml_tensor * masks = ggml_view_3d(ctx0, model.dspark_mask_embd, n_embd, block_drafts - 1, 1,
+ model.dspark_mask_embd->nb[1], model.dspark_mask_embd->nb[2], 0);
+ masks = ggml_repeat_4d(ctx0, ggml_cast(ctx0, masks, GGML_TYPE_F32), n_embd, block_drafts - 1, n_blocks, 1);
+ inpL = ggml_concat(ctx0, inpL, masks, 1);
+ }
+ inpL = ggml_reshape_2d(ctx0, inpL, n_embd, n_tokens);
+ } else {
+ // not a draft block (e.g. graph reserve)
+ inpL = ggml_get_rows(ctx0, tok_embd, inp_tokens);
+ }
+ cb(inpL, "inp_noise_embd", -1);
+
+ // the whole block shares the mask row of its first query (the anchor)
+ inp_attn->self_attn_mask_swa.block_uniform = true;
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ {
+ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq, cur);
+ cb(q, "q", il);
+
+ // {n_embd_head_nope, n_tokens, n_head}
+ ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_nope, n_head, n_tokens,
+ ggml_row_size(q->type, n_embd_head_k_mla),
+ ggml_row_size(q->type, n_embd_head_k_mla) * n_head, 0);
+ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+
+ // {kv_lora_rank, n_head, n_tokens}
+ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
+ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
+ cb(q_nope_absorbed, "q_nope_absorbed", il);
+
+ ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_rot, n_head, n_tokens,
+ ggml_row_size(q->type, n_embd_head_k_mla),
+ ggml_row_size(q->type, n_embd_head_k_mla) * n_head,
+ ggml_row_size(q->type, n_embd_head_nope));
+ q_pe = rope(q_pe);
+ cb(q_pe, "q_pe", il);
+
+ ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
+ cb(Qcur, "Qcur", il);
+
+ auto [Kcur, Vcur] = build_kv(cur, il);
+ cb(Kcur, "Kcur", il);
+
+ // apply wv_b after attn, so the K rotation is already undone
+ ggml_tensor * out = build_attn(inp_attn,
+ nullptr, nullptr, nullptr,
+ Qcur, Kcur, Vcur,
+ nullptr, nullptr, nullptr,
+ kq_scale, il);
+ cb(out, "attn_latent", il);
+
+ // {kv_lora_rank, n_tokens, n_head} -> {n_embd_head_v_mla, n_tokens, n_head}
+ out = ggml_reshape_3d(ctx0, out, kv_lora_rank, n_head, n_tokens);
+ out = ggml_mul_mat(ctx0, layer.wv_b, ggml_permute(ctx0, out, 0, 2, 1, 3));
+ out = ggml_cont_2d(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3), n_embd_head_v_mla*n_head, n_tokens);
+
+ cur = build_lora_mm(layer.wo, out, layer.wo_s);
+ cb(cur, "attn_out", il);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
+ cb(ffn_inp, "ffn_inp", il);
+
+ cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ cur = build_ffn(cur,
+ layer.ffn_up, nullptr, nullptr,
+ layer.ffn_gate, nullptr, nullptr,
+ layer.ffn_down, nullptr, nullptr,
+ nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+
+ inpL = ggml_add(ctx0, cur, ffn_inp);
+ cb(inpL, "l_out", il);
+ }
+
+ ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ const auto [output, output_s] = get_lm_head(*this, model);
+ ggml_tensor * logits = build_lora_mm(output, cur, output_s);
+
+ // softcap does not apply to the markov bias
+ if (hparams.f_final_logit_softcapping) {
+ logits = ggml_scale(ctx0, logits, 1.0f / hparams.f_final_logit_softcapping);
+ logits = ggml_tanh(ctx0, logits);
+ logits = ggml_scale(ctx0, logits, hparams.f_final_logit_softcapping);
+ }
+ cb(logits, "result_output", -1);
+ res->t_logits = logits;
+
+ if (model.dspark_markov_w1) {
+ build_dspark_markov_head(*this, model, inp_tokens, cur);
+ }
+}
diff --git a/src/models/lilicorr.cpp b/src/models/lilicorr.cpp
new file mode 100644
index 000000000..0c5527a65
--- /dev/null
+++ b/src/models/lilicorr.cpp
@@ -0,0 +1,280 @@
+#include "models.h"
+
+void llama_model_dflash::load_hparams_lilicorr(llama_model_loader & ml) {
+ ml.get_key(LLM_KV_LILICORR_CANDIDATE_TOPK, hparams.lilicorr_candidate_topk);
+ ml.get_key(LLM_KV_LILICORR_NUM_HEADS, hparams.lilicorr_num_heads);
+ ml.get_key(LLM_KV_LILICORR_NUM_LAYERS, hparams.lilicorr_layers);
+ ml.get_key(LLM_KV_LILICORR_VECTOR_EPS, hparams.lilicorr_vector_eps);
+ ml.get_key(LLM_KV_LILICORR_LOGIT_SCALE, hparams.lilicorr_logit_scale);
+ if (hparams.dflash_block_size < 2 || hparams.lilicorr_hidden_size == 0 ||
+ hparams.lilicorr_num_heads == 0 || hparams.lilicorr_hidden_size % hparams.lilicorr_num_heads ||
+ hparams.lilicorr_candidate_topk == 0 || hparams.lilicorr_layers == 0 ||
+ hparams.dflash_selector_top_k > 0 ||
+ (uint64_t) hparams.lilicorr_candidate_topk * (hparams.lilicorr_candidate_topk + 1) > hparams.n_embd) {
+ throw std::runtime_error("invalid LiLiCorr DFlash metadata");
+ }
+}
+
+void llama_model_dflash::load_tensors_lilicorr() {
+ LLAMA_LOAD_LOCALS;
+
+ const int64_t hidden = hparams.lilicorr_hidden_size;
+ const int64_t top_k = hparams.lilicorr_candidate_topk;
+ const int64_t heads = hparams.lilicorr_num_heads;
+ uint32_t factor = 0, mlp = 0;
+ ml->get_key(LLM_KV_LILICORR_FACTOR_DIM, factor);
+ ml->get_key(LLM_KV_LILICORR_MLP_HIDDEN_SIZE, mlp);
+ const int64_t slots = hparams.dflash_block_size - 1;
+ const int feature_bid = n_layer_all + hparams.lilicorr_layers;
+ if (factor == 0 || mlp == 0 || hparams.lilicorr_layers > layers.size()) {
+ throw std::runtime_error("invalid LiLiCorr head dimensions");
+ }
+ if (hidden != n_embd) {
+ lilicorr_token_proj = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "weight"), { n_embd, hidden }, 0);
+ lilicorr_token_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "bias"), { hidden }, 0);
+ lilicorr_token_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
+ }
+ lilicorr_pass_hidden_proj = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "weight"), { hparams.n_embd, hidden }, 0);
+ lilicorr_pass_hidden_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "bias"), { hidden }, 0);
+ lilicorr_pass_hidden_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
+ lilicorr_feature_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", feature_bid), { 5 }, 0);
+ lilicorr_feature_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", feature_bid), { 5 }, 0);
+ lilicorr_feature_mlp_1 = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", feature_bid), { 5, hidden }, 0);
+ lilicorr_feature_mlp_1_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", feature_bid), { hidden }, 0);
+ lilicorr_feature_mlp_3 = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", feature_bid), { hidden, hidden }, 0);
+ lilicorr_feature_mlp_3_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", feature_bid), { hidden }, 0);
+ lilicorr_slot_embedding = create_tensor(tn(LLM_TENSOR_LILICORR_SLOT_EMBD, "weight"), { hidden, slots }, 0);
+ lilicorr_rank_embedding = create_tensor(tn(LLM_TENSOR_LILICORR_RANK_EMBD, "weight"), { hidden, top_k }, 0);
+ lilicorr_attention_bias = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_BIAS, "bias"), { slots * top_k, slots * top_k, heads }, 0);
+ lilicorr_context_proj = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "weight"), { n_embd, hidden }, 0);
+ lilicorr_context_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "bias"), { hidden }, 0);
+ lilicorr_context_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
+ for (uint32_t i = 0; i < hparams.lilicorr_layers; ++i) {
+ const int bid = n_layer_all + i;
+ layers[i].lilicorr_attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", bid), { hidden }, 0);
+ layers[i].lilicorr_mlp_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", bid), { hidden }, 0);
+ layers[i].lilicorr_attn_in_proj = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "weight", i), { hidden, 3 * hidden }, 0);
+ layers[i].lilicorr_attn_in_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "bias", i), { 3 * hidden }, 0);
+ layers[i].lilicorr_attn_out_proj = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", bid), { hidden, hidden }, 0);
+ layers[i].lilicorr_attn_out_proj_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", bid), { hidden }, 0);
+ layers[i].lilicorr_attn_out_proj_s = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
+ layers[i].lilicorr_mlp_0 = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", bid), { hidden, mlp }, 0);
+ layers[i].lilicorr_mlp_0_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", bid), { mlp }, 0);
+ layers[i].lilicorr_mlp_0_s = create_tensor(tn(LLM_TENSOR_FFN_UP, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
+ layers[i].lilicorr_mlp_2 = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", bid), { mlp, hidden }, 0);
+ layers[i].lilicorr_mlp_2_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", bid), { hidden }, 0);
+ layers[i].lilicorr_mlp_2_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED);
+ }
+ lilicorr_output_norm = create_tensor(tn(LLM_TENSOR_LILICORR_OUTPUT_NORM, "weight"), { hidden }, 0);
+ lilicorr_anchor_norm = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_NORM, "weight"), { hidden }, 0);
+ lilicorr_factor_input_proj = create_tensor(tn(LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "weight"), { 3 * hidden, hidden }, 0);
+ lilicorr_factor_input_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "bias"), { hidden }, 0);
+ lilicorr_out_head = create_tensor(tn(LLM_TENSOR_LILICORR_OUT_HEAD, "weight"), { hidden, factor }, 0);
+ lilicorr_out_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_OUT_HEAD, "bias"), { factor }, 0);
+ lilicorr_in_head = create_tensor(tn(LLM_TENSOR_LILICORR_IN_HEAD, "weight"), { hidden, factor }, 0);
+ lilicorr_in_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_IN_HEAD, "bias"), { factor }, 0);
+ lilicorr_anchor_out_head = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "weight"), { hidden, factor }, 0);
+ lilicorr_anchor_out_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "bias"), { factor }, 0);
+
+ load_tensors_dflash2_conv();
+}
+
+static void build_lilicorr_scores(llm_graph_context & g, const llama_model & model,
+ ggml_tensor * tok_embd, ggml_tensor * anchors) {
+ ggml_context * ctx0 = g.ctx0;
+ auto & res = g.res;
+
+ const auto & hparams = g.hparams;
+ const int64_t n_tokens = g.n_tokens;
+ const int64_t n_embd = g.n_embd;
+
+ const int64_t hidden = hparams.lilicorr_hidden_size;
+ const int64_t top_k = hparams.lilicorr_candidate_topk;
+ const int64_t heads = hparams.lilicorr_num_heads;
+ const int64_t n_blocks = g.ubatch.n_seqs_unq;
+ GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
+ GGML_ASSERT(res->t_logits->ne[1] == n_tokens);
+ GGML_ASSERT(anchors && tok_embd);
+
+ const int64_t tokens_per_block = n_tokens / n_blocks;
+ const int64_t block_size = hparams.dflash_block_size;
+ const int64_t row_used = top_k + top_k * top_k;
+ const int64_t slots = block_size - 1;
+ const int64_t nodes = slots * top_k;
+ const int64_t factor_dim = model.lilicorr_out_head->ne[1];
+ GGML_ASSERT(tokens_per_block == block_size);
+
+ ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k);
+ ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens);
+ ggml_tensor * candidate_logits = ggml_reshape_2d(ctx0,
+ ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens);
+ // Rank features and embeddings require candidates in descending logit order.
+ ggml_tensor * order = ggml_argsort(ctx0, candidate_logits, GGML_SORT_ORDER_DESC);
+ ggml_tensor * candidate_rows = ggml_reshape_3d(ctx0, candidates, 1, top_k, n_tokens);
+ candidates = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, candidate_rows, order), top_k, n_tokens);
+ ggml_tensor * prob_rows = ggml_reshape_3d(ctx0, ggml_soft_max(ctx0, res->t_logits), 1, res->t_logits->ne[0], n_tokens);
+ ggml_tensor * probs = ggml_reshape_2d(ctx0,
+ ggml_get_rows(ctx0, prob_rows, candidates), top_k, n_tokens);
+
+ ggml_tensor * cand_run = ggml_reshape_3d(ctx0, candidates, top_k, n_tokens / n_blocks, n_blocks);
+ ggml_tensor * probs_run = ggml_reshape_3d(ctx0, probs, top_k, n_tokens / n_blocks, n_blocks);
+ cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_run, top_k, slots, n_blocks,
+ cand_run->nb[1], cand_run->nb[2], cand_run->nb[1]));
+ probs_run = ggml_cont(ctx0, ggml_view_3d(ctx0, probs_run, top_k, slots, n_blocks,
+ probs_run->nb[1], probs_run->nb[2], probs_run->nb[1]));
+
+ ggml_tensor * lp = ggml_log(ctx0, probs_run);
+ ggml_tensor * lp_blk = ggml_reshape_3d(ctx0, lp, top_k, slots, n_blocks);
+ ggml_tensor * top_lp = ggml_view_3d(ctx0, lp_blk, 1, slots, n_blocks,
+ lp_blk->nb[1], lp_blk->nb[2], 0);
+ ggml_tensor * relative = ggml_sub(ctx0, lp_blk, top_lp);
+ ggml_tensor * feature = ggml_reshape_4d(ctx0, lp_blk, 1, top_k, slots, n_blocks);
+
+ ggml_tensor * rank = ggml_arange(ctx0, 0, top_k, 1);
+ rank = ggml_repeat(ctx0, ggml_reshape_3d(ctx0, rank, top_k, 1, 1), lp_blk);
+ ggml_tensor * first = ggml_scale_bias(ctx0, ggml_step(ctx0, rank), -1.0f, 1.0f);
+
+ for (ggml_tensor * item : { probs_run, relative,
+ ggml_scale(ctx0, rank, top_k == 1 ? 0.0f : 1.0f / (top_k - 1)), first }) {
+ feature = ggml_concat(ctx0, feature, ggml_reshape_4d(ctx0, item, 1, top_k, slots, n_blocks), 0);
+ }
+ feature = ggml_reshape_2d(ctx0, feature, 5, nodes * n_blocks);
+
+ ggml_tensor * embeds = ggml_get_rows(ctx0, tok_embd, ggml_reshape_1d(ctx0, cand_run, nodes * n_blocks));
+ ggml_tensor * token_states = n_embd == hidden ? embeds : ggml_add(ctx0, g.build_lora_mm(model.lilicorr_token_proj, embeds, model.lilicorr_token_proj_s), model.lilicorr_token_proj_b);
+
+ ggml_tensor * pass_blk = ggml_reshape_3d(ctx0, res->t_embd, n_embd, block_size, n_blocks);
+ ggml_tensor * passes = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_view_3d(ctx0, pass_blk, n_embd, slots, n_blocks,
+ pass_blk->nb[1], pass_blk->nb[2], pass_blk->nb[1])),
+ n_embd, slots * n_blocks);
+
+ ggml_tensor * pass_states = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_pass_hidden_proj, passes, model.lilicorr_pass_hidden_proj_s), model.lilicorr_pass_hidden_proj_b);
+
+ pass_states = ggml_reshape_4d(ctx0, pass_states, hidden, 1, slots, n_blocks);
+ pass_states = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, pass_states, hidden, top_k, slots, n_blocks), hidden, nodes * n_blocks);
+
+ ggml_tensor * feature_states = ggml_norm(ctx0, feature, 1.0e-5f);
+ feature_states = ggml_add(ctx0, ggml_mul(ctx0, feature_states, model.lilicorr_feature_norm), model.lilicorr_feature_norm_b);
+ feature_states = ggml_silu(ctx0, ggml_add(ctx0, g.build_lora_mm(model.lilicorr_feature_mlp_1, feature_states), model.lilicorr_feature_mlp_1_b));
+ feature_states = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_feature_mlp_3, feature_states), model.lilicorr_feature_mlp_3_b);
+
+ ggml_tensor * slot = ggml_reshape_4d(ctx0, model.lilicorr_slot_embedding, hidden, 1, slots, 1);
+ ggml_tensor * rank_embd = ggml_reshape_4d(ctx0, model.lilicorr_rank_embedding, hidden, top_k, 1, 1);
+ ggml_tensor * pos = ggml_add(ctx0, ggml_repeat_4d(ctx0, slot, hidden, top_k, slots, n_blocks), ggml_repeat_4d(ctx0, rank_embd, hidden, top_k, slots, n_blocks));
+ pos = ggml_reshape_2d(ctx0, pos, hidden, nodes * n_blocks);
+
+ ggml_tensor * x = ggml_add(ctx0, ggml_add(ctx0, token_states, pass_states), feature_states);
+ x = ggml_add(ctx0, x, pos);
+ ggml_tensor * anchor_blk = ggml_reshape_3d(ctx0, anchors, n_embd, block_size, n_blocks);
+ ggml_tensor * anchor = ggml_reshape_2d(ctx0,
+ ggml_cont(ctx0, ggml_view_3d(ctx0, anchor_blk, n_embd, 1, n_blocks,
+ anchor_blk->nb[1], anchor_blk->nb[2], 0)), n_embd, n_blocks);
+ ggml_tensor * anchor_state = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_context_proj, anchor, model.lilicorr_context_proj_s), model.lilicorr_context_proj_b);
+
+ const int64_t head_dim = hidden / heads;
+ for (uint32_t layer = 0; layer < hparams.lilicorr_layers; ++layer) {
+ const auto & weights = model.layers[layer];
+ ggml_tensor * attn_input = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), weights.lilicorr_attn_norm);
+ ggml_tensor * qkv = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_attn_in_proj, attn_input), weights.lilicorr_attn_in_proj_b);
+ qkv = ggml_reshape_3d(ctx0, qkv, 3 * hidden, nodes, n_blocks);
+ ggml_tensor * q = ggml_cont(ctx0, ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
+ qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], 0));
+ ggml_tensor * k = ggml_cont(ctx0, ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
+ qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], hidden * qkv->nb[0]));
+ ggml_tensor * v = ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks,
+ qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], 2 * hidden * qkv->nb[0]);
+ ggml_tensor * attn = ggml_soft_max_ext(ctx0, ggml_mul_mat(ctx0, k, q), model.lilicorr_attention_bias,
+ 1.0f / std::sqrt(static_cast<float>(head_dim)), 0.0f);
+ ggml_tensor * joined = ggml_out_prod(ctx0, v, ggml_transpose(ctx0, attn));
+ joined = ggml_permute(ctx0, joined, 0, 2, 1, 3);
+ ggml_tensor * attn_output = ggml_cont_2d(ctx0, joined, hidden, nodes * n_blocks);
+ attn_output = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_attn_out_proj, attn_output, weights.lilicorr_attn_out_proj_s), weights.lilicorr_attn_out_proj_b);
+ x = ggml_add(ctx0, x, attn_output);
+ ggml_tensor * mlp = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), weights.lilicorr_mlp_norm);
+ mlp = ggml_silu(ctx0, ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_mlp_0, mlp, weights.lilicorr_mlp_0_s), weights.lilicorr_mlp_0_b));
+ mlp = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_mlp_2, mlp, weights.lilicorr_mlp_2_s), weights.lilicorr_mlp_2_b);
+ x = ggml_add(ctx0, x, mlp);
+ }
+
+ x = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), model.lilicorr_output_norm);
+ anchor_state = ggml_mul(ctx0, ggml_rms_norm(ctx0, anchor_state, hparams.f_norm_rms_eps), model.lilicorr_anchor_norm);
+ ggml_tensor * anchor_nodes = ggml_reshape_3d(ctx0, anchor_state, hidden, 1, n_blocks);
+ anchor_nodes = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, anchor_nodes, hidden, nodes, n_blocks, 1), hidden, nodes * n_blocks);
+ ggml_tensor * factor_in = ggml_concat(ctx0, x, anchor_nodes, 0);
+ factor_in = ggml_concat(ctx0, factor_in, ggml_mul(ctx0, x, anchor_nodes), 0);
+ ggml_tensor * factors = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_factor_input_proj, factor_in), model.lilicorr_factor_input_proj_b);
+ factors = ggml_silu(ctx0, factors);
+
+ ggml_tensor * out = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_out_head, factors), model.lilicorr_out_head_b);
+ out = ggml_reshape_4d(ctx0, ggml_l2_norm(ctx0, out, hparams.lilicorr_vector_eps), factor_dim, top_k, slots, n_blocks);
+
+ ggml_tensor * in = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_in_head, factors), model.lilicorr_in_head_b);
+ in = ggml_reshape_4d(ctx0, ggml_l2_norm(ctx0, in, hparams.lilicorr_vector_eps), factor_dim, top_k, slots, n_blocks);
+
+ ggml_tensor * anchor_out = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_anchor_out_head, anchor_state), model.lilicorr_anchor_out_head_b);
+ anchor_out = ggml_l2_norm(ctx0, anchor_out, hparams.lilicorr_vector_eps);
+ anchor_out = ggml_reshape_3d(ctx0, anchor_out, factor_dim, 1, n_blocks);
+
+ ggml_tensor * first_in = ggml_cont(ctx0, ggml_view_4d(ctx0, in, factor_dim, top_k, 1, n_blocks,
+ in->nb[1], in->nb[2], in->nb[3], 0));
+ ggml_tensor * first_score = ggml_scale(ctx0,
+ ggml_mul_mat(ctx0, first_in, ggml_reshape_4d(ctx0, anchor_out, factor_dim, 1, 1, n_blocks)),
+ hparams.lilicorr_logit_scale);
+ ggml_tensor * scores = ggml_pad(ctx0,
+ ggml_reshape_3d(ctx0, first_score, top_k, 1, n_blocks), top_k * top_k - top_k, 0, 0, 0);
+ if (slots > 1) {
+ ggml_tensor * next = ggml_cont(ctx0, ggml_view_4d(ctx0, in, factor_dim, top_k, slots - 1, n_blocks,
+ in->nb[1], in->nb[2], in->nb[3], in->nb[2]));
+ ggml_tensor * previous = ggml_cont(ctx0, ggml_view_4d(ctx0, out, factor_dim, top_k, slots - 1, n_blocks,
+ out->nb[1], out->nb[2], out->nb[3], 0));
+ ggml_tensor * score = ggml_scale(ctx0, ggml_mul_mat(ctx0, next, previous), hparams.lilicorr_logit_scale);
+ scores = ggml_concat(ctx0, scores,
+ ggml_reshape_3d(ctx0, score, top_k * top_k, slots - 1, n_blocks), 1);
+ }
+
+ ggml_tensor * row = ggml_concat(ctx0, ggml_cast(ctx0, ggml_reshape_3d(ctx0, cand_run, top_k, slots, n_blocks),
+ GGML_TYPE_F32), scores, 0);
+ ggml_tensor * packed = ggml_pad_ext(ctx0, row, 0, n_embd - row_used, 1, 0, 0, 0, 0, 0);
+ packed = ggml_reshape_2d(ctx0, packed, n_embd, n_tokens);
+ g.cb(packed, "lilicorr_lattice", -1);
+ res->t_h_nextn = packed;
+ ggml_build_forward_expand(g.gf, packed);
+}
+
+llama_model_dflash::graph_lilicorr::graph_lilicorr(const llm_graph_params & params) : graph_base(params) {
+ ASSERT_EMBD_OR_TOKEN(ubatch);
+ if (ubatch.embd || ubatch.embd_row) {
+ // Expose context features for retrieval by LiliCorr
+ res->t_h_nextn = build_inject();
+ return;
+ }
+
+ // only a full block per sequence has the anchors the scorer needs
+ const bool is_block = ubatch.embd_state && n_tokens == (int64_t) hparams.dflash_block_size * ubatch.n_seqs_unq;
+
+ ggml_tensor * inp_tokens = nullptr;
+ ggml_tensor * inp_anchor = nullptr;
+ if (is_block) {
+ inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp_tokens);
+ res->t_inp_tokens = inp_tokens;
+
+ // inp_anchor passes the context features to the lattice scorer
+ auto inp = std::make_unique<llm_graph_input_embd_h>(n_embd, n_embd);
+ inp->tokens = inp_tokens;
+ inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
+ ggml_set_input(inp->h);
+ inp_anchor = inp->h;
+ res->add_input(std::move(inp));
+ } else {
+ inp_tokens = build_inp_tokens();
+ }
+
+ ggml_tensor * cur = build_layers(inp_tokens);
+
+ build_logits(cur);
+
+ if (is_block) {
+ build_lilicorr_scores(*this, model, get_tok_embd(*this, model), inp_anchor);
+ }
+}
diff --git a/src/models/models.h b/src/models/models.h
index cfef363b6..a978c6fab 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -1522,6 +1522,34 @@ struct llama_model_dflash : public llama_model_base {
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
+ enum variant_type {
+ VARIANT_DFLASH,
+ VARIANT_DFLASH2,
+ VARIANT_DSPARK,
+ VARIANT_LILICORR,
+ };
+
+ enum backbone_type {
+ BACKBONE_BASE,
+ BACKBONE_DSV4,
+ BACKBONE_DSV41,
+ BACKBONE_MISTRAL_LARGE4,
+ };
+
+ variant_type variant = VARIANT_DFLASH;
+ backbone_type backbone = BACKBONE_BASE;
+
+ void load_hparams_lilicorr(llama_model_loader & ml);
+ void load_hparams_dsv4(llama_model_loader & ml);
+ void load_hparams_mistral_large4(llama_model_loader & ml);
+ void load_tensors_dflash2();
+ void load_tensors_dflash2_conv();
+ void load_tensors_dspark(int64_t n_vocab_draft);
+ void load_tensors_lilicorr();
+ void load_tensors_layers_base();
+ void load_tensors_layers_dsv4();
+ void load_tensors_layers_mistral_large4();
+
template <bool is_enc>
struct graph : public llm_graph_context {
graph(const llm_graph_params & params);
@@ -1529,6 +1557,42 @@ struct llama_model_dflash : public llama_model_base {
ggml_tensor * build_inp_embd_enc() const;
};
+ // inputs and stages shared by the decoder graphs of the heads
+ struct graph_base : public llm_graph_context {
+ int64_t n_embd_head;
+ bool use_iswa;
+ int sections[4];
+
+ ggml_tensor * inp_pos;
+
+ llm_graph_input_attn_kv * inp_attn = nullptr;
+ llm_graph_input_attn_kv_iswa * inp_attn_iswa = nullptr;
+
+ graph_base(const llm_graph_params & params);
+
+ ggml_tensor * build_inject();
+ ggml_tensor * build_inp_tokens();
+ ggml_tensor * build_layers(ggml_tensor * inp_tokens);
+ ggml_tensor * build_logits(ggml_tensor * cur);
+ ggml_tensor * build_rope(ggml_tensor * cur, ggml_tensor * pos);
+ };
+
+ struct graph_dflash : public graph_base {
+ graph_dflash(const llm_graph_params & params);
+ };
+
+ struct graph_dflash2 : public graph_base {
+ graph_dflash2(const llm_graph_params & params);
+ };
+
+ struct graph_dspark : public graph_base {
+ graph_dspark(const llm_graph_params & params);
+ };
+
+ struct graph_lilicorr : public graph_base {
+ graph_lilicorr(const llm_graph_params & params);
+ };
+
struct graph_dsv4 : public llama_model_deepseek4::graph {
graph_dsv4(const llm_graph_params & params);
};
@@ -1543,6 +1607,14 @@ struct llama_model_dflash : public llama_model_base {
graph_mistral_large4(const llm_graph_params & params);
};
+ // graph helpers shared by the backbones
+ static ggml_tensor * build_inject_inp(llm_graph_context & g, const llama_model & model);
+ static ggml_tensor * get_tok_embd(const llm_graph_context & g, const llama_model & model);
+ static std::pair<ggml_tensor *, ggml_tensor *> get_lm_head(const llm_graph_context & g, const llama_model & model);
+
+ // DFlash2 local conv, run by the base layers of DFlash2 and LiliCorr drafts
+ static ggml_tensor * build_dflash2_conv(llm_graph_context & g, ggml_tensor * hidden, ggml_tensor * dynamic, ggml_tensor * base, int side);
+
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
diff --git a/tools/cli/README.md b/tools/cli/README.md
index 475becd2e..1dc83ee81 100644
--- a/tools/cli/README.md
+++ b/tools/cli/README.md
@@ -211,7 +211,7 @@
| `--spec-draft-device, -devd, --device-draft <dev1,dev2,..>` | comma-separated list of devices to use for offloading the draft model (none = don't offload, default: follows --device)<br/>use --list-devices to see a list of available devices |
| `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)<br/>(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) |
| `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)<br/>(env: LLAMA_ARG_SPEC_DRAFT_MODEL) |
-| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)<br/><br/>(env: LLAMA_ARG_SPEC_TYPE) |
+| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,draft-dflash2,draft-lilicorr,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)<br/><br/>(env: LLAMA_ARG_SPEC_TYPE) |
| `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) |
| `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) |
| `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) |
diff --git a/tools/server/README.md b/tools/server/README.md
index 066ecb4dd..9f5cd62cc 100644
--- a/tools/server/README.md
+++ b/tools/server/README.md
@@ -277,7 +277,7 @@ For the full list of features, please refer to [server's changelog](https://gith
| `--spec-draft-device, -devd, --device-draft <dev1,dev2,..>` | comma-separated list of devices to use for offloading the draft model (none = don't offload, default: follows --device)<br/>use --list-devices to see a list of available devices |
| `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)<br/>(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) |
| `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)<br/>(env: LLAMA_ARG_SPEC_DRAFT_MODEL) |
-| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)<br/><br/>(env: LLAMA_ARG_SPEC_TYPE) |
+| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,draft-dflash2,draft-lilicorr,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)<br/><br/>(env: LLAMA_ARG_SPEC_TYPE) |
| `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) |
| `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) |
| `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) |