diff --git a/common/sampling.cpp b/common/sampling.cpp
@@ -4,6 +4,7 @@
#include "fit.h"
#include "log.h"
#include "reasoning-budget.h"
+#include "speculative.h"
#include "ggml.h"
@@ -12,6 +13,7 @@
#include <climits>
#include <cmath>
#include <cstring>
+#include <random>
#include <unordered_map>
#include <vector>
@@ -121,10 +123,14 @@ struct common_sampler {
llama_token_data_array cur_p;
+ uint32_t speculative_seed;
+ std::mt19937 speculative_rng;
+
void reset() {
prev.clear();
llama_sampler_reset(chain);
+ speculative_rng.seed(speculative_seed);
}
void set_logits(struct llama_context * ctx, int idx) {
@@ -424,6 +430,8 @@ struct common_sampler * common_sampler_init(
params.backend_sampling = false;
}
+ // Keep verifier randomness independent from both target and draft sampling.
+ const uint32_t speculative_seed = llama_sampler_get_seed(chain) ^ 0x9e3779b9U;
auto * result = new common_sampler {
/* .params = */ params,
/* .grmr = */ grmr,
@@ -432,6 +440,8 @@ struct common_sampler * common_sampler_init(
/* .prev = */ ring_buffer<llama_token>(std::max(32, params.n_prev)),
/* .cur = */ {},
/* .cur_p = */ {},
+ /* .speculative_seed = */ speculative_seed,
+ /* .speculative_rng = */ std::mt19937(speculative_seed),
};
return result;
@@ -515,6 +525,8 @@ struct common_sampler * common_sampler_clone(common_sampler * gsmpl) {
/* .prev = */ gsmpl->prev,
/* .cur = */ gsmpl->cur,
/* .cur_p = */ gsmpl->cur_p,
+ /* .speculative_seed = */ gsmpl->speculative_seed,
+ /* .speculative_rng = */ gsmpl->speculative_rng,
};
}
@@ -535,6 +547,8 @@ void common_sampler_copy(const common_sampler * src, common_sampler * dst) {
dst->cur = src->cur;
dst->cur_p = src->cur_p;
dst->cur_p.data = src->cur_p.data ? dst->cur.data() : nullptr; // re-point to dst's buffer
+ dst->speculative_seed = src->speculative_seed;
+ dst->speculative_rng = src->speculative_rng;
dst->t_total_us = src->t_total_us;
}
@@ -705,6 +719,79 @@ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sample
return result;
}
+std::vector<llama_token> common_sampler_sample_and_accept_n(
+ struct common_sampler * gsmpl,
+ struct llama_context * ctx,
+ const std::vector<int> & idxs,
+ const llama_tokens & draft,
+ const std::vector<common_speculative_token_dist> & dists,
+ bool grammar_first) {
+ GGML_ASSERT(idxs.size() == draft.size() + 1);
+ GGML_ASSERT(dists.size() == draft.size());
+
+ std::vector<llama_token> result;
+ result.reserve(idxs.size());
+
+ std::uniform_real_distribution<float> uniform(0.0f, 1.0f);
+ size_t i = 0;
+ for (; i < draft.size(); ++i) {
+ // Residual sampling needs the target distribution after every constraint.
+ const llama_token fallback = common_sampler_sample(gsmpl, ctx, idxs[i], true);
+ const auto & q = dists[i];
+ GGML_ASSERT(q.ids.size() == q.probs.size());
+
+ std::unordered_map<llama_token, float> q_probs;
+ q_probs.reserve(q.ids.size());
+ for (size_t j = 0; j < q.ids.size(); ++j) {
+ q_probs[q.ids[j]] += q.probs[j];
+ }
+ const auto q_prob = [&](llama_token id) {
+ const auto it = q_probs.find(id);
+ return it == q_probs.end() ? 0.0f : it->second;
+ };
+
+ auto * p = common_sampler_get_candidates(gsmpl, false);
+ float p_draft = 0.0f;
+ const float q_draft = q_prob(draft[i]);
+ for (size_t j = 0; j < p->size; ++j) {
+ if (p->data[j].id == draft[i]) {
+ p_draft = p->data[j].p;
+ break;
+ }
+ }
+
+ if (q_draft > 0.0f && uniform(gsmpl->speculative_rng) * q_draft <= p_draft) {
+ common_sampler_accept(gsmpl, draft[i], true);
+ result.push_back(draft[i]);
+ continue;
+ }
+
+ std::vector<float> residual(p->size);
+ float residual_sum = 0.0f;
+ for (size_t j = 0; j < p->size; ++j) {
+ residual[j] = std::max(0.0f, p->data[j].p - q_prob(p->data[j].id));
+ residual_sum += residual[j];
+ }
+
+ llama_token id = fallback;
+ if (residual_sum > 0.0f) {
+ std::discrete_distribution<size_t> sample(residual.begin(), residual.end());
+ id = p->data[sample(gsmpl->speculative_rng)].id;
+ }
+ common_sampler_accept(gsmpl, id, true);
+ result.push_back(id);
+ break;
+ }
+
+ if (i == draft.size()) {
+ const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first);
+ common_sampler_accept(gsmpl, id, true);
+ result.push_back(id);
+ }
+
+ return result;
+}
+
std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) {
std::vector<int> idxs(draft.size() + 1);
for (size_t i = 0; i < idxs.size(); ++i) {
@@ -714,6 +801,19 @@ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sample
return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first);
}
+std::vector<llama_token> common_sampler_sample_and_accept_n(
+ struct common_sampler * gsmpl,
+ struct llama_context * ctx,
+ const llama_tokens & draft,
+ const std::vector<common_speculative_token_dist> & dists,
+ bool grammar_first) {
+ std::vector<int> idxs(draft.size() + 1);
+ for (size_t i = 0; i < idxs.size(); ++i) {
+ idxs[i] = i;
+ }
+ return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, dists, grammar_first);
+}
+
uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl) {
return llama_sampler_get_seed(gsmpl->chain);
}
diff --git a/common/sampling.h b/common/sampling.h
@@ -33,6 +33,7 @@
//
struct common_sampler;
+struct common_speculative_token_dist;
// llama_sampler API overloads
@@ -85,6 +86,23 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co
//
std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first = false);
+// maximal-coupling verification for stochastic speculative decoding
+std::vector<llama_token> common_sampler_sample_and_accept_n(
+ struct common_sampler * gsmpl,
+ struct llama_context * ctx,
+ const std::vector<int> & idxs,
+ const llama_tokens & draft,
+ const std::vector<common_speculative_token_dist> & dists,
+ bool grammar_first = false);
+
+// assume idxs == [ 0, 1, 2, ..., draft.size() ]
+std::vector<llama_token> common_sampler_sample_and_accept_n(
+ struct common_sampler * gsmpl,
+ struct llama_context * ctx,
+ const llama_tokens & draft,
+ const std::vector<common_speculative_token_dist> & dists,
+ bool grammar_first = false);
+
// assume idxs == [ 0, 1, 2, ..., draft.size() ]
std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false);
diff --git a/common/speculative.cpp b/common/speculative.cpp
@@ -14,10 +14,12 @@
#include <algorithm>
#include <cassert>
+#include <chrono>
+#include <cinttypes>
#include <cstring>
#include <iomanip>
#include <map>
-#include <cinttypes>
+#include <random>
#define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
#define SPC_TRC(fmt, ...) LOG_TRC("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
@@ -923,6 +925,11 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
int32_t block_size = 0;
llama_token mask_token_id = 0;
+ bool is_dflash2 = false;
+ int32_t selector_top_k = 0;
+ std::vector<std::mt19937> selector_rng;
+ std::vector<bool> selector_reset;
+
// draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
const bool is_dspark;
@@ -966,6 +973,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
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.selector_top_k", buf, sizeof(buf)) >= 0) {
+ selector_top_k = std::atoi(buf);
+ is_dflash2 = selector_top_k > 0;
+ }
}
mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
@@ -991,14 +1002,17 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
for (auto & s : smpls) {
common_params_sampling sparams;
sparams.no_perf = false;
- sparams.top_k = 10;
+ sparams.top_k = is_dflash2 ? selector_top_k : 10;
sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K };
s.reset(common_sampler_init(model_dft, sparams));
}
+ selector_rng.resize(n_seq);
+ selector_reset.assign(n_seq, true);
+
// offload draft sampling to the backend
backend_chains.assign(n_seq, nullptr);
- if (this->params.backend_sampling) {
+ if (this->params.backend_sampling && !is_dflash2) {
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));
@@ -1017,7 +1031,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
}
- llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
+ // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
+ llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2);
llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention
}
@@ -1048,6 +1063,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
return;
}
+ selector_reset[seq_id] = true;
+
const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(params.ctx_dft), seq_id);
if (pos_max < N - 1) {
LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - process() did not run on every prefill ubatch. "
@@ -1074,87 +1091,68 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
const int32_t n_tokens = batch_in.n_tokens;
- // per-seq inclusive batch range (assumes each seq's tokens are contiguous in the batch)
- std::vector<int32_t> i_batch_beg(n_seq, -1);
- std::vector<int32_t> i_batch_end(n_seq, -1);
- for (int32_t k = 0; k < n_tokens; ++k) {
- GGML_ASSERT(batch_in.n_seq_id[k] == 1);
- const llama_seq_id seq_id = batch_in.seq_id[k][0];
- if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
- continue;
- }
- i_batch_end[seq_id] = k;
- if (i_batch_beg[seq_id] < 0) {
- i_batch_beg[seq_id] = k;
- }
- }
-
auto * ctx_tgt = this->params.ctx_tgt;
auto * ctx_dft = this->params.ctx_dft;
const int32_t n_ubatch = (int32_t) llama_n_ubatch(ctx_dft);
- for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
- if (i_batch_beg[seq_id] < 0) {
- continue;
- }
- const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;
-
- for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
- const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
-
- // gather this chunk's target features, interleaved by extract layer
- features_buf.resize((size_t) n_chunk * n_embd_enc);
- for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
- const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
- if (!layer) {
- GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
- }
- for (int32_t i = 0; i < n_chunk; ++i) {
- float * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
- const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;
- std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
- }
+ // Flatten token-wise encoder work into shared chunks while preserving each row's position and sequence.
+ for (int32_t offset = 0; offset < n_tokens; offset += n_ubatch) {
+ const int32_t n_chunk = std::min(n_ubatch, n_tokens - offset);
+ features_buf.resize((size_t) n_chunk * n_embd_enc);
+ for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
+ const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
+ if (!layer) {
+ GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
+ }
+ for (int32_t i = 0; i < n_chunk; ++i) {
+ float * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
+ const float * src = layer + (size_t) (offset + i) * n_embd_tgt;
+ std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
}
+ }
- // fuse extracted features through DFlash encoder
- llama_batch enc_batch = {
- /*.n_tokens =*/ n_chunk,
- /*.token =*/ nullptr,
- /*.embd =*/ features_buf.data(),
- /*.pos =*/ nullptr,
- /*.n_seq_id =*/ nullptr,
- /*.seq_id =*/ nullptr,
- /*.logits =*/ nullptr,
- };
+ llama_batch enc_batch = {
+ /*.n_tokens =*/ n_chunk,
+ /*.token =*/ nullptr,
+ /*.embd =*/ features_buf.data(),
+ /*.pos =*/ nullptr,
+ /*.n_seq_id =*/ nullptr,
+ /*.seq_id =*/ nullptr,
+ /*.logits =*/ nullptr,
+ };
- int32_t rc = llama_encode(ctx_dft, enc_batch);
- if (rc != 0) {
- LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
- __func__, rc, (int) n_chunk, (int) offset);
- return false;
- }
+ int32_t rc = llama_encode(ctx_dft, enc_batch);
+ if (rc != 0) {
+ LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+ __func__, rc, (int) n_chunk, (int) offset);
+ return false;
+ }
- const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
- GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
+ const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
+ GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
- // inject the DFlash decoder K/V cache at the tokens' target positions
- batch_inject.n_tokens = n_chunk;
- std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
+ batch_inject.n_tokens = n_chunk;
+ std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
+ for (int32_t i = 0; i < n_chunk; ++i) {
+ const int32_t j = offset + i;
+ GGML_ASSERT(batch_in.n_seq_id[j] == 1);
+ const llama_seq_id seq_id = batch_in.seq_id[j][0];
+ GGML_ASSERT(seq_id >= 0 && seq_id < (llama_seq_id) n_seq);
+ batch_inject.pos[i] = batch_in.pos[j];
+ batch_inject.n_seq_id[i] = 1;
+ batch_inject.seq_id[i][0] = seq_id;
+ batch_inject.logits[i] = false;
+ }
- for (int32_t i = 0; i < n_chunk; ++i) {
- batch_inject.pos[i] = batch_in.pos[i_batch_beg[seq_id] + offset + i];
- batch_inject.n_seq_id[i] = 1;
- batch_inject.seq_id[i][0] = seq_id;
- batch_inject.logits[i] = false;
- }
- rc = llama_decode(ctx_dft, batch_inject);
- if (rc != 0) {
- LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
- __func__, rc, (int) n_chunk, (int) offset);
- return false;
- }
+ rc = llama_decode(ctx_dft, batch_inject);
+ if (rc != 0) {
+ LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+ __func__, rc, (int) n_chunk, (int) offset);
+ return false;
}
+ // The server may switch contexts before the next draft decode.
+ llama_synchronize(ctx_dft);
}
return true;
@@ -1186,7 +1184,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
i_block_beg[seq_id] = batch.n_tokens;
n_block [seq_id] = n_block_tokens;
for (int32_t i = 0; i < n_block_tokens; ++i) {
- common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true);
+ common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, !is_dflash2);
}
}
@@ -1214,6 +1212,63 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
auto & result = *dp.result;
+ if (dp.dists) {
+ dp.dists->clear();
+ }
+
+ if (is_dflash2) {
+ GGML_ASSERT(dp.temperature <= 0.0f || dp.dists);
+ const float * lattice = llama_get_embeddings_nextn(ctx_dft);
+ GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
+
+ if (selector_reset[seq_id]) {
+ uint32_t seed = dp.seed;
+ if (seed == LLAMA_DEFAULT_SEED) {
+ seed = (uint32_t) std::chrono::high_resolution_clock::now().time_since_epoch().count();
+ }
+ selector_rng[seq_id].seed(seed ^ 0x85ebca6bU);
+ selector_reset[seq_id] = false;
+ }
+
+ 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 float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k;
+
+ if (dp.temperature > 0.0f) {
+ common_speculative_token_dist dist;
+ dist.ids.resize(selector_top_k);
+ dist.probs.resize(selector_top_k);
+ const float max_score = *std::max_element(scores, scores + selector_top_k);
+ float sum = 0.0f;
+ for (int32_t k = 0; k < selector_top_k; ++k) {
+ dist.ids[k] = (llama_token) row[k];
+ dist.probs[k] = std::exp((scores[k] - max_score) / dp.temperature);
+ sum += dist.probs[k];
+ }
+ for (float & p : dist.probs) {
+ p /= sum;
+ }
+ std::discrete_distribution<int32_t> sample(dist.probs.begin(), dist.probs.end());
+ predecessor = sample(selector_rng[seq_id]);
+ result.push_back(dist.ids[predecessor]);
+ dp.dists->push_back(std::move(dist));
+ } else {
+ predecessor = (int32_t) std::distance(scores,
+ std::max_element(scores, scores + selector_top_k));
+ result.push_back((llama_token) row[predecessor]);
+ }
+ }
+
+ if (result.size() < (size_t) params.n_min) {
+ result.clear();
+ if (dp.dists) {
+ dp.dists->clear();
+ }
+ }
+ continue;
+ }
+
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;
@@ -2339,7 +2394,7 @@ common_params common_base_params_to_speculative(const common_params & params) {
result.n_outputs_max = params.n_parallel;
result.n_outputs_max_per_seq = 1;
- // dflash/dspark decode the whole noise block in a single pass and sample every block position on the backend
+ // dflash/dspark decode every sequence's full noise block in one pass
// TODO: refactor such properties to be announced by the speculative types
// something like `struct common_speculative_type_props common_speculative_type_get_props(...);`
const bool has_block_draft = std::any_of(
@@ -2351,6 +2406,8 @@ common_params common_base_params_to_speculative(const common_params & params) {
// per-seq output positions: DFlash decodes anchor + n_max masks (n_max + 1); DSpark n_max -> +1 covers both
const int32_t per_seq = std::max(1, params_spec.n_max + 1);
result.n_outputs_max = params.n_parallel * per_seq;
+ result.n_batch = std::max(result.n_batch, result.n_outputs_max);
+ result.n_ubatch = std::max(result.n_ubatch, result.n_outputs_max);
if (params_spec.backend_sampling) {
result.n_outputs_max_per_seq = per_seq;
}
diff --git a/common/speculative.h b/common/speculative.h
@@ -5,6 +5,11 @@
struct common_speculative;
+struct common_speculative_token_dist {
+ llama_tokens ids;
+ std::vector<float> probs;
+};
+
// comma separated list the provided types
std::string common_speculative_type_name_str(const std::vector<enum common_speculative_type> & types);
@@ -60,6 +65,12 @@ struct common_speculative_draft_params {
// the generated draft from the last _draft() call
llama_tokens * result;
+
+ // optional sparse proposal distributions, one per draft token
+ std::vector<common_speculative_token_dist> * dists = nullptr;
+
+ float temperature = 0.0f;
+ uint32_t seed = LLAMA_DEFAULT_SEED;
};
common_speculative_draft_params & common_speculative_get_draft_params(common_speculative * spec, llama_seq_id seq_id);
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
@@ -334,6 +334,11 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_TARGET_LAYERS, "%s.target_layers" },
{ LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" },
+ { LLM_KV_DFLASH_BLOCK_SIZE, "%s.block_size" },
+ { LLM_KV_DFLASH_CONV_KERNEL_SIZE, "%s.conv_kernel_size" },
+ { LLM_KV_DFLASH_CONV_GROUP_SIZE, "%s.conv_group_size" },
+ { LLM_KV_DFLASH_SELECTOR_RANK, "%s.selector_rank" },
+ { LLM_KV_DFLASH_SELECTOR_TOP_K, "%s.selector_top_k" },
{ LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" },
{ LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" },
@@ -644,6 +649,13 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" },
{ LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" },
{ LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" },
+ { LLM_TENSOR_DFLASH_ATTN_CONV_BASE, "blk.%d.attn_conv_base" },
+ { LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "blk.%d.attn_conv_proj" },
+ { LLM_TENSOR_DFLASH_FFN_CONV_BASE, "blk.%d.ffn_conv_base" },
+ { LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "blk.%d.ffn_conv_proj" },
+ { LLM_TENSOR_DFLASH_SELECTOR_PREV, "selector_predecessor" },
+ { LLM_TENSOR_DFLASH_SELECTOR_NEXT, "selector_successor" },
+ { LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "selector_hidden" },
};
// declare information about the model weight tensors:
@@ -909,6 +921,13 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
{LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
{LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_DFLASH_ATTN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+ {LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_DFLASH_FFN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+ {LLM_TENSOR_DFLASH_FFN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_DFLASH_SELECTOR_PREV, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
+ {LLM_TENSOR_DFLASH_SELECTOR_NEXT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
+ {LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
};
LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {}
diff --git a/src/llama-arch.h b/src/llama-arch.h
@@ -380,6 +380,11 @@ enum llm_kv {
LLM_KV_TARGET_LAYERS,
LLM_KV_TARGET_HIDDEN_SIZE,
+ LLM_KV_DFLASH_BLOCK_SIZE,
+ LLM_KV_DFLASH_CONV_KERNEL_SIZE,
+ LLM_KV_DFLASH_CONV_GROUP_SIZE,
+ LLM_KV_DFLASH_SELECTOR_RANK,
+ LLM_KV_DFLASH_SELECTOR_TOP_K,
LLM_KV_NORM_BEFORE_RESIDUAL,
LLM_KV_NORM_BEFORE_FC,
@@ -652,6 +657,13 @@ enum llm_tensor {
LLM_TENSOR_DSPARK_MARKOV_W1,
LLM_TENSOR_DSPARK_MARKOV_W2,
LLM_TENSOR_DSPARK_CONF_PROJ,
+ LLM_TENSOR_DFLASH_ATTN_CONV_BASE,
+ LLM_TENSOR_DFLASH_ATTN_CONV_PROJ,
+ LLM_TENSOR_DFLASH_FFN_CONV_BASE,
+ LLM_TENSOR_DFLASH_FFN_CONV_PROJ,
+ LLM_TENSOR_DFLASH_SELECTOR_PREV,
+ LLM_TENSOR_DFLASH_SELECTOR_NEXT,
+ LLM_TENSOR_DFLASH_SELECTOR_HIDDEN,
};
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
@@ -2314,6 +2314,12 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
}
}
+ if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) {
+ const uint32_t selector_tokens = std::min<uint32_t>(
+ n_tokens, model.hparams.dflash_block_size * cparams.n_seq_max);
+ res += 32*selector_tokens;
+ }
+
uint32_t n_sampling_nodes = 0;
uint32_t n_sampling_nodes_max = 0;
for (const auto & [seq_id, sampler] : sampling.samplers) {
diff --git a/src/llama-graph.h b/src/llama-graph.h
@@ -1326,6 +1326,8 @@ struct llm_graph_context {
void build_sampling() const;
+ virtual void build_post_sampling() const {}
+
//
// dense (out)
//
diff --git a/src/llama-hparams.h b/src/llama-hparams.h
@@ -214,6 +214,12 @@ struct llama_hparams {
// output embedding dimension (0 = use n_embd)
uint32_t n_embd_out_impl = 0;
+ uint32_t dflash_block_size = 0;
+ uint32_t dflash_conv_kernel_size = 0;
+ uint32_t dflash_conv_group_size = 0;
+ uint32_t dflash_selector_rank = 0;
+ uint32_t dflash_selector_top_k = 0;
+
// llama4 smallthinker
uint32_t n_moe_layer_step = 0;
uint32_t n_no_rope_layer_step = 4;
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
@@ -2459,6 +2459,7 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
// add backend sampling layers (if any)
llm->build_sampling();
+ llm->build_post_sampling();
// if the gguf model was converted with --sentence-transformers-dense-modules
// there will be two additional dense projection layers
diff --git a/src/llama-model.h b/src/llama-model.h
@@ -362,6 +362,11 @@ struct llama_layer {
struct ggml_tensor * ffn_exp_probs_b = nullptr;
struct ggml_tensor * ffn_gate_tid2eid = nullptr;
+ struct ggml_tensor * dflash_attn_conv_base = nullptr;
+ struct ggml_tensor * dflash_attn_conv_proj = nullptr;
+ struct ggml_tensor * dflash_ffn_conv_base = nullptr;
+ struct ggml_tensor * dflash_ffn_conv_proj = nullptr;
+
// mamba proj
struct ggml_tensor * ssm_in = nullptr;
struct ggml_tensor * ssm_x = nullptr;
@@ -648,6 +653,10 @@ struct llama_model {
struct ggml_tensor * dspark_conf_proj = nullptr;
struct ggml_tensor * dspark_conf_proj_b = nullptr;
+ struct ggml_tensor * dflash_selector_prev = nullptr;
+ struct ggml_tensor * dflash_selector_next = nullptr;
+ struct ggml_tensor * dflash_selector_hidden = nullptr;
+
// unified vector to store target-model extracted layer ids in eagle3, dflash, etc.
std::vector<int32_t> target_layer_ids;
diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
@@ -7,6 +7,16 @@
void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+ ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, false);
+ hparams.f_final_logit_softcapping = 0.0f;
+ ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
+ ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false);
+
+ ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false);
+ ml.get_key(LLM_KV_DFLASH_CONV_KERNEL_SIZE, hparams.dflash_conv_kernel_size, false);
+ 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);
if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) {
throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata");
@@ -112,6 +122,29 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
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");
+ 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);
+ }
+
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
@@ -187,6 +220,16 @@ 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 (selector_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);
+ }
}
}
@@ -222,7 +265,7 @@ ggml_tensor * llama_model_dflash::graph<true>::build_inp_embd_enc() const {
// DFlash Encoder: processes target model features through feature fusion layer
template <>
-llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model) {
ggml_tensor * cur = build_inp_embd_enc();
cur = build_lora_mm(model.fc, cur, model.fc_s);
@@ -345,11 +388,69 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
ggml_build_forward_expand(g.gf, out);
}
+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;
+ hidden = ggml_cont_2d(ctx0, hidden, hidden_size, 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 * grouped = ggml_reshape_3d(ctx0, hidden, group_size, n_groups, n_tokens);
+ 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 * result = nullptr;
+ for (int64_t tap = 0; tap < kernel_size; ++tap) {
+ ggml_tensor * values = blocks;
+ if (tap > 0) {
+ ggml_tensor * zeros = ggml_fill(ctx0,
+ ggml_new_tensor_3d(ctx0, hidden->type, hidden_size, std::min(tap, block_size), n_blocks), 0.0f);
+ if (tap < block_size) {
+ ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks,
+ blocks->nb[1], blocks->nb[2], 0);
+ values = ggml_concat(ctx0, zeros, previous, 1);
+ } else {
+ values = zeros;
+ }
+ }
+ values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens);
+
+ ggml_tensor * coeff = ggml_view_2d(ctx0, coeffs_side, n_groups, n_tokens,
+ coeffs_side->nb[2], tap * coeffs_side->nb[1]);
+ coeff = ggml_cont(ctx0, coeff);
+ coeff = ggml_reshape_3d(ctx0, coeff, 1, n_groups, n_tokens);
+ coeff = ggml_reshape_2d(ctx0, ggml_repeat(ctx0, coeff, grouped), hidden_size, n_tokens);
+
+ ggml_tensor * base_tap = ggml_view_1d(ctx0, base, hidden_size,
+ tap * base->nb[1] + side * base->nb[2]);
+ ggml_tensor * weight = ggml_add(ctx0, coeff, ggml_repeat(ctx0, base_tap, hidden));
+ ggml_tensor * term = ggml_mul(ctx0, weight, values);
+ result = result ? ggml_add(ctx0, result, term) : term;
+ }
+ return result;
+}
+
// 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 llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
@@ -449,10 +550,14 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
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_tokens = inp->tokens;
ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
+ if (hparams.f_embedding_scale != 0.0f) {
+ inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
+ }
cb(inpL, "inp_noise_embd", -1);
res->add_input(std::move(inp));
@@ -463,6 +568,13 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il);
cb(noise_norm, "noise_norm", il);
+ ggml_tensor * attn_dynamic = nullptr;
+ if (layer.dflash_attn_conv_proj) {
+ attn_dynamic = build_lora_mm(layer.dflash_attn_conv_proj, noise_norm);
+ noise_norm = build_dflash2_conv(*this, noise_norm, attn_dynamic, layer.dflash_attn_conv_base, 0);
+ cb(noise_norm, "attn_conv_in", il);
+ }
+
ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm);
ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm);
ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm);
@@ -493,12 +605,24 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il)
: build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ if (attn_dynamic) {
+ cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);
+ cb(cur, "attn_conv_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, NULL, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
+ ggml_tensor * ffn_dynamic = nullptr;
+ if (layer.dflash_ffn_conv_proj) {
+ ffn_dynamic = build_lora_mm(layer.dflash_ffn_conv_proj, cur);
+ cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 0);
+ cb(cur, "ffn_conv_in", il);
+ }
+
cur = build_ffn(cur,
layer.ffn_up, NULL, layer.ffn_up_s,
layer.ffn_gate, NULL, layer.ffn_gate_s,
@@ -507,6 +631,11 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(cur, "ffn_out", il);
+ if (ffn_dynamic) {
+ cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 1);
+ cb(cur, "ffn_conv_out", il);
+ }
+
cur = ggml_add(ctx0, cur, ffn_inp);
cb(cur, "l_out", il);
@@ -531,6 +660,15 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
cur = build_lora_mm(output, cur, output_s);
+ if (hparams.f_logit_scale != 0.0f) {
+ cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);
+ }
+ if (hparams.f_final_logit_softcapping > 0.0f) {
+ cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
+ cur = ggml_tanh(ctx0, cur);
+ cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
+ }
+
// reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t
if (model.d2t) {
const int64_t n_draft_vocab = cur->ne[0];
@@ -557,6 +695,98 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
}
}
+template <bool is_enc>
+void llama_model_dflash::graph<is_enc>::build_post_sampling() const {
+ if constexpr (is_enc) {
+ return;
+ }
+
+ if (!model.dflash_selector_hidden || !res->t_logits) {
+ return;
+ }
+
+ const int64_t top_k = hparams.dflash_selector_top_k;
+ const int64_t rank = hparams.dflash_selector_rank;
+ const int64_t n_blocks = ubatch.n_seqs_unq;
+ GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
+ GGML_ASSERT(res->t_logits->ne[1] == n_tokens);
+ ggml_tensor * tokens = res->get_inp_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);
+ 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);
+
+ std::vector<ggml_tensor *> candidate_ids(block_size);
+ std::vector<ggml_tensor *> unary_logits(block_size);
+ for (int64_t pos = 1; pos < block_size; ++pos) {
+ candidate_ids[pos] = ggml_cont_2d(ctx0,
+ ggml_view_2d(ctx0, candidates, top_k, n_blocks,
+ tokens_per_block * candidates->nb[1], pos * candidates->nb[1]),
+ top_k, n_blocks);
+ unary_logits[pos] = ggml_cont_2d(ctx0,
+ ggml_view_2d(ctx0, unary, top_k, n_blocks,
+ tokens_per_block * unary->nb[1], pos * unary->nb[1]),
+ top_k, n_blocks);
+ }
+
+ ggml_tensor * hidden = build_lora_mm(model.dflash_selector_hidden, res->t_embd);
+
+ ggml_tensor * anchor_ids = ggml_view_2d(ctx0, tokens, 1, n_blocks,
+ tokens_per_block * tokens->nb[0], 0);
+ anchor_ids = ggml_cont_1d(ctx0, anchor_ids, n_blocks);
+
+ ggml_tensor * packed = ggml_fill(ctx0,
+ ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f);
+
+ for (int64_t pos = 1; pos < block_size; ++pos) {
+ ggml_tensor * ids = candidate_ids[pos];
+ ggml_tensor * unary = unary_logits[pos];
+ ggml_tensor * successor = ggml_get_rows(ctx0, model.dflash_selector_next,
+ ggml_reshape_1d(ctx0, ids, top_k * n_blocks));
+ successor = ggml_reshape_3d(ctx0, successor, rank, top_k, n_blocks);
+
+ ggml_tensor * hidden_pos = ggml_cont(ctx0, ggml_view_2d(ctx0, hidden, rank, n_blocks,
+ tokens_per_block * hidden->nb[1], pos * hidden->nb[1]));
+ hidden_pos = ggml_reshape_3d(ctx0, hidden_pos, rank, 1, n_blocks);
+
+ ggml_tensor * predecessor;
+ if (pos == 1) {
+ predecessor = ggml_get_rows(ctx0, model.dflash_selector_prev, anchor_ids);
+ predecessor = ggml_reshape_3d(ctx0, predecessor, rank, 1, n_blocks);
+ } else {
+ predecessor = ggml_get_rows(ctx0, model.dflash_selector_prev,
+ ggml_reshape_1d(ctx0, candidate_ids[pos - 1], top_k * n_blocks));
+ predecessor = ggml_reshape_3d(ctx0, predecessor, rank, top_k, n_blocks);
+ }
+
+ ggml_tensor * conditioned = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, hidden_pos, predecessor));
+ ggml_tensor * scores = ggml_mul_mat(ctx0, successor, conditioned);
+ if (pos == 1) {
+ scores = ggml_repeat_4d(ctx0, scores, top_k, top_k, n_blocks, 1);
+ }
+ ggml_tensor * unary_3d = ggml_reshape_3d(ctx0, unary, top_k, 1, n_blocks);
+ scores = ggml_add(ctx0, scores, ggml_repeat(ctx0, unary_3d, scores));
+
+ ggml_tensor * row = ggml_concat(ctx0,
+ ggml_cast(ctx0, ids, GGML_TYPE_F32),
+ ggml_reshape_2d(ctx0, scores, top_k * top_k, n_blocks), 0);
+ row = ggml_pad(ctx0, row, n_embd - row->ne[0], 0, 0, 0);
+ row = ggml_reshape_3d(ctx0, row, n_embd, 1, n_blocks);
+ packed = ggml_concat(ctx0, packed, row, 1);
+ }
+
+ packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks);
+ cb(packed, "dflash2_lattice", -1);
+ res->t_h_nextn = packed;
+ ggml_build_forward_expand(gf, packed);
+}
+
// 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
diff --git a/src/models/models.h b/src/models/models.h
@@ -1342,9 +1342,12 @@ struct llama_model_dflash : public llama_model_base {
template <bool is_enc>
struct graph : public llm_graph_context {
+ const llama_model & model;
+
graph(const llama_model & model, const llm_graph_params & params);
ggml_tensor * build_inp_embd_enc() const;
+ void build_post_sampling() const override;
};
struct graph_dsv4 : public llama_model_deepseek4::graph {