From 5ecbe1ac17ec0484c5b44af0bd580cdc9c428ed4 Mon Sep 17 00:00:00 2001 From: Jian Chen Date: Tue, 18 Aug 2026 02:08:02 -0700 Subject: [PATCH 01/16] support DFlash2 --- common/sampling.cpp | 100 ++++++++ common/sampling.h | 18 ++ common/speculative.cpp | 207 ++++++++++------ common/speculative.h | 11 + conversion/__init__.py | 1 + conversion/qwen.py | 34 ++- .../speculative-simple/speculative-simple.cpp | 16 +- gguf-py/gguf/constants.py | 25 ++ gguf-py/gguf/gguf_writer.py | 12 + gguf-py/gguf/tensor_mapping.py | 28 +++ src/llama-arch.cpp | 19 ++ src/llama-arch.h | 12 + src/llama-context.cpp | 6 + src/llama-graph.h | 2 + src/llama-hparams.h | 6 + src/llama-model.cpp | 1 + src/llama-model.h | 9 + src/models/dflash.cpp | 234 +++++++++++++++++- src/models/models.h | 3 + tools/server/server-context.cpp | 15 +- 20 files changed, 676 insertions(+), 83 deletions(-) diff --git a/common/sampling.cpp b/common/sampling.cpp index 06dea1e1cce..da8bdc1b112 100644 --- 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 #include #include +#include #include #include @@ -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(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 common_sampler_sample_and_accept_n(struct common_sample return result; } +std::vector common_sampler_sample_and_accept_n( + struct common_sampler * gsmpl, + struct llama_context * ctx, + const std::vector & idxs, + const llama_tokens & draft, + const std::vector & dists, + bool grammar_first) { + GGML_ASSERT(idxs.size() == draft.size() + 1); + GGML_ASSERT(dists.size() == draft.size()); + + std::vector result; + result.reserve(idxs.size()); + + std::uniform_real_distribution 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 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 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 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 common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) { std::vector idxs(draft.size() + 1); for (size_t i = 0; i < idxs.size(); ++i) { @@ -714,6 +801,19 @@ std::vector common_sampler_sample_and_accept_n(struct common_sample return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first); } +std::vector common_sampler_sample_and_accept_n( + struct common_sampler * gsmpl, + struct llama_context * ctx, + const llama_tokens & draft, + const std::vector & dists, + bool grammar_first) { + std::vector 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 index ced3c8364b3..5f3273eae4b 100644 --- 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 common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first = false); +// maximal-coupling verification for stochastic speculative decoding +std::vector common_sampler_sample_and_accept_n( + struct common_sampler * gsmpl, + struct llama_context * ctx, + const std::vector & idxs, + const llama_tokens & draft, + const std::vector & dists, + bool grammar_first = false); + +// assume idxs == [ 0, 1, 2, ..., draft.size() ] +std::vector common_sampler_sample_and_accept_n( + struct common_sampler * gsmpl, + struct llama_context * ctx, + const llama_tokens & draft, + const std::vector & dists, + bool grammar_first = false); + // assume idxs == [ 0, 1, 2, ..., draft.size() ] std::vector 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 index ae55e357d51..b69ab19515c 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -14,10 +14,12 @@ #include #include +#include +#include #include #include #include -#include +#include #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 selector_rng; + std::vector 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 i_batch_beg(n_seq, -1); - std::vector 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 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 index 12ae31b7de5..f7dd6526737 100644 --- 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 probs; +}; + // comma separated list the provided types std::string common_speculative_type_name_str(const std::vector & 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 * 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/conversion/__init__.py b/conversion/__init__.py index 3232a1050bf..f6897330575 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -54,6 +54,7 @@ "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", + "DFlash2DraftModel": "qwen", "Qwen3DSparkModel": "qwen", "DSparkDraftModel": "qwen", "DSparkSpeculator": "qwen", diff --git a/conversion/qwen.py b/conversion/qwen.py index 26b10452b6f..ab2a0fb4dce 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -639,7 +639,7 @@ class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): model_arch = gguf.MODEL_ARCH.QWEN35MOE -@ModelBase.register("DFlashDraftModel") +@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel") @ModelBase.example("z-lab/Qwen3.5-9B-DFlash") class DFlashModel(Qwen3Model): model_arch = gguf.MODEL_ARCH.DFLASH @@ -678,9 +678,31 @@ def set_vocab(self): def set_gguf_parameters(self): super().set_gguf_parameters() - block_size = self.hparams.get("block_size", 16) - self.gguf_writer.add_block_size(block_size) dflash_config = self.hparams.get("dflash_config", {}) + block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16)) + self.gguf_writer.add_block_size(block_size) + + if "conv_kernel_size" in dflash_config: + self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"])) + self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"])) + self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"])) + self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"])) + + output_multiplier = dflash_config.get( + "output_multiplier", self.hparams.get("output_multiplier") + ) + if output_multiplier is not None: + self.gguf_writer.add_logit_scale(float(output_multiplier)) + softcap = dflash_config.get( + "final_logit_softcapping", self.hparams.get("final_logit_softcapping") + ) + if softcap is not None and float(softcap) > 0: + self.gguf_writer.add_final_logit_softcapping(float(softcap)) + embedding_scale = dflash_config.get( + "input_embedding_scale", self.hparams.get("input_embedding_scale") + ) + if embedding_scale is not None: + self.gguf_writer.add_embedding_scale(float(embedding_scale)) target_layer_ids = dflash_config.get("target_layer_ids", []) if target_layer_ids: @@ -706,6 +728,12 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True): return + if name in ( + "model.candidate_selector.predecessor_codebook", + "model.candidate_selector.successor_codebook", + ): + name += ".weight" + yield from super().modify_tensors(data_torch, name, bid) diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index 487ae03abfa..d074f270f97 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -66,7 +66,8 @@ int main(int argc, char ** argv) { llama_context * ctx_dft = params.speculative.draft.ctx_dft; // check if the context supports partial sequence removal - const bool use_ckpt_tgt = common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; + const auto ctx_tgt_seq_rm_type = common_context_can_seq_rm(ctx_tgt); + const bool use_ckpt_tgt = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; const bool use_ckpt_dft = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; if (use_ckpt_tgt) { @@ -152,6 +153,7 @@ int main(int argc, char ** argv) { llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1); llama_tokens draft; + std::vector dists; common_prompt_checkpoint ckpt; @@ -192,6 +194,9 @@ int main(int argc, char ** argv) { /* .id_last = */ id_last, /* .prompt = */ &prompt_tgt, /* .result = */ &draft, // output + /* .dists = */ &dists, + /* .temperature = */ params.sampling.temp, + /* .seed = */ common_sampler_get_seed(smpl.get()), }; common_speculative_draft(spec); @@ -255,7 +260,12 @@ int main(int argc, char ** argv) { // available logits from the batch and sample the next token until we run out of logits or the sampler // disagrees with the draft // - auto ids = common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft); + const bool can_rollback = + ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART || + (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() <= llama_n_rs_seq(ctx_tgt)); + auto ids = can_rollback && params.sampling.temp > 0.0f && dists.size() == draft.size() + ? common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft, dists) + : common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft); //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str()); @@ -268,6 +278,7 @@ int main(int argc, char ** argv) { LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, n_draft); draft = std::move(ids); + dists.clear(); { ckpt.load_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); @@ -325,6 +336,7 @@ int main(int argc, char ** argv) { // clear the draft since it has been consumed draft.clear(); + dists.clear(); { LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past); diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d043c9b6ecc..e3dc7a7a892 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -162,6 +162,10 @@ class LLM: TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" + CONV_KERNEL_SIZE = "{arch}.conv_kernel_size" + CONV_GROUP_SIZE = "{arch}.conv_group_size" + SELECTOR_RANK = "{arch}.selector_rank" + SELECTOR_TOP_K = "{arch}.selector_top_k" SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" @@ -1133,6 +1137,13 @@ class MODEL_TENSOR(IntEnum): DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed DSPARK_MARKOV_W2 = auto() # markov head: bias projection DSPARK_CONF_PROJ = auto() # confidence head + DFLASH_ATTN_CONV_BASE = auto() + DFLASH_ATTN_CONV_PROJ = auto() + DFLASH_FFN_CONV_BASE = auto() + DFLASH_FFN_CONV_PROJ = auto() + DFLASH_SELECTOR_PREV = auto() + DFLASH_SELECTOR_NEXT = auto() + DFLASH_SELECTOR_HIDDEN = auto() # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -1875,6 +1886,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base", + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj", + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base", + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj", + MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor", + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor", + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden", MODEL_TENSOR.D2T: "d2t", } @@ -4865,6 +4883,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1, MODEL_TENSOR.DSPARK_MARKOV_W2, MODEL_TENSOR.DSPARK_CONF_PROJ, + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE, + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ, + MODEL_TENSOR.DFLASH_FFN_CONV_BASE, + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ, + MODEL_TENSOR.DFLASH_SELECTOR_PREV, + MODEL_TENSOR.DFLASH_SELECTOR_NEXT, + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN, ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 9e0914fd866..5d65ec6fea1 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -981,6 +981,18 @@ def add_sliding_window(self, value: int) -> None: def add_block_size(self, value: int) -> None: self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value) + def add_conv_kernel_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.CONV_KERNEL_SIZE.format(arch=self.arch), value) + + def add_conv_group_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.CONV_GROUP_SIZE.format(arch=self.arch), value) + + def add_selector_rank(self, value: int) -> None: + self.add_uint32(Keys.LLM.SELECTOR_RANK.format(arch=self.arch), value) + + def add_selector_top_k(self, value: int) -> None: + self.add_uint32(Keys.LLM.SELECTOR_TOP_K.format(arch=self.arch), value) + def add_sample_from_anchor(self, value: bool) -> None: self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 3292942b413..0802d2c04f3 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1353,6 +1353,34 @@ class TensorNameMap: "model.confidence_head.proj", # dspark ), + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: ( + "model.layers.{bid}.attention_conv.base_kernel", + ), + + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: ( + "model.layers.{bid}.attention_conv.kernel_projection", + ), + + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: ( + "model.layers.{bid}.mlp_conv.base_kernel", + ), + + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: ( + "model.layers.{bid}.mlp_conv.kernel_projection", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_PREV: ( + "model.candidate_selector.predecessor_codebook", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: ( + "model.candidate_selector.successor_codebook", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: ( + "model.candidate_selector.hidden_projection", + ), + MODEL_TENSOR.CLS: ( "classifier", # jina "classifier.dense", # roberta diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5b88bde14d1..8fd4c8974fb 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -334,6 +334,11 @@ static const std::map 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_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_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 index 8042120a25f..7e44842b163 100644 --- 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 index 52f8d53672a..51148b2adcc 100644 --- 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( + 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 index 94324c7457e..605733f32b0 100644 --- 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 index e91ce1cc3c7..5ddcae84dd1 100644 --- 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 index 0d74a2135b6..e31f38ea413 100644 --- 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 index 4412ef08e74..e2cd04946e5 100644 --- 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 target_layer_ids; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 5b70a517949..776d34dd2c6 100644 --- 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::build_inp_embd_enc() const { // DFlash Encoder: processes target model features through feature fusion layer template <> -llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { +llama_model_dflash::graph::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::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { +llama_model_dflash::graph::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::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::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::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::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::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::graph(const llama_model & model, const llm_gra } } +template +void llama_model_dflash::graph::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(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 candidate_ids(block_size); + std::vector 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 index 180b30a46d7..969bec894c7 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1342,9 +1342,12 @@ struct llama_model_dflash : public llama_model_base { template 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 { diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 842e4203cd2..df11f31b8af 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -209,6 +209,7 @@ struct server_slot { common_speculative * spec; llama_tokens spec_draft; + std::vector spec_dists; llama_tokens spec_prompt; std::vector spec_i_batch; common_prompt_checkpoint spec_ckpt; @@ -337,6 +338,7 @@ struct server_slot { if (can_speculate()) { spec_draft.clear(); + spec_dists.clear(); spec_i_batch.clear(); spec_ckpt.clear(); } @@ -2938,6 +2940,9 @@ struct server_context_impl { /* .id_last = */ slot.sampled, /* .prompt = */ &slot.spec_prompt, /* .result = */ &slot.spec_draft, + /* .dists = */ &slot.spec_dists, + /* .temperature = */ slot.task->params.sampling.temp, + /* .seed = */ common_sampler_get_seed(slot.smpl.get()), }; drafting.push_back(&slot); @@ -3817,7 +3822,13 @@ struct server_context_impl { common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); - auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); + const bool can_rollback = + ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART || + (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && n_draft <= llama_n_rs_seq(ctx_tgt)); + auto accepted = can_rollback && slot.task->params.sampling.temp > 0.0f && + slot.spec_dists.size() == slot.spec_draft.size() + ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, slot.spec_dists) + : common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1); @@ -3838,6 +3849,7 @@ struct server_context_impl { // partial acceptance is not supported by the context -> truncate the draft and restore the state slot.spec_is_replay = true; slot.spec_draft = std::move(accepted); + slot.spec_dists.clear(); const auto & ckpt = slot.spec_ckpt; @@ -3865,6 +3877,7 @@ struct server_context_impl { common_speculative_accept(spec.get(), slot.id, accepted.size() - 1); slot.spec_draft = std::move(accepted); + slot.spec_dists.clear(); } const auto ids = std::move(slot.spec_draft); From 1deefcca395743049c3820ab8f9b15043f3e9446 Mon Sep 17 00:00:00 2001 From: SubSir Date: Fri, 21 Aug 2026 02:01:25 -0700 Subject: [PATCH 02/16] Add p_min in DFlash2 Assisted-by: Claude Opus 5 --- common/speculative.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/common/speculative.cpp b/common/speculative.cpp index b69ab19515c..01c7ee41b87 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1251,11 +1251,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } std::discrete_distribution sample(dist.probs.begin(), dist.probs.end()); predecessor = sample(selector_rng[seq_id]); + if (dist.probs[predecessor] < params.p_min) { + break; + } 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)); + 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 < selector_top_k; ++k) { + sum += std::exp(scores[k] - scores[predecessor]); + } + if (1.0f / sum < params.p_min) { + break; + } + } result.push_back((llama_token) row[predecessor]); } } From c314bb10d2b747e105fe122f0aefb080263d4e01 Mon Sep 17 00:00:00 2001 From: SubSir Date: Fri, 21 Aug 2026 16:39:16 -0700 Subject: [PATCH 03/16] Revert unnecessary changes Assisted-by: Claude Opus 5 --- common/speculative.cpp | 125 +++++++++++++++++++++++------------------ 1 file changed, 71 insertions(+), 54 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 01c7ee41b87..113578c7aa2 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1002,7 +1002,7 @@ 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 = is_dflash2 ? selector_top_k : 10; + sparams.top_k = 10; sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K }; s.reset(common_sampler_init(model_dft, sparams)); } @@ -1091,68 +1091,87 @@ 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 i_batch_beg(n_seq, -1); + std::vector 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); - // 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)); - } + 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; - 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, - }; + for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) { + const int32_t n_chunk = std::min(n_ubatch, n_rows - offset); - 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; - } + // 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)); + } + } - const float * inp_g = llama_get_embeddings_nextn(ctx_dft); - GGML_ASSERT(inp_g && "DFlash encoder produced no output."); + // 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, + }; - 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; - } + 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; + } - 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; + 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)); + + 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; + } } - // The server may switch contexts before the next draft decode. - llama_synchronize(ctx_dft); } return true; @@ -2407,7 +2426,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 every sequence's full noise block in one pass + // dflash/dspark decode the whole noise block in a single pass and sample every block position on the backend // 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( @@ -2419,8 +2438,6 @@ 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; } From b8a4a8532b17283e230b6732dc1735a5b6b7937c Mon Sep 17 00:00:00 2001 From: SubSir Date: Fri, 21 Aug 2026 16:48:54 -0700 Subject: [PATCH 04/16] Revert draft sampling in rejection sampling Assisted-by: Claude Opus 5 --- common/sampling.cpp | 100 ------------------ common/sampling.h | 18 ---- common/speculative.cpp | 66 ++---------- common/speculative.h | 11 -- .../speculative-simple/speculative-simple.cpp | 16 +-- tools/server/server-context.cpp | 15 +-- 6 files changed, 12 insertions(+), 214 deletions(-) diff --git a/common/sampling.cpp b/common/sampling.cpp index da8bdc1b112..06dea1e1cce 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -4,7 +4,6 @@ #include "fit.h" #include "log.h" #include "reasoning-budget.h" -#include "speculative.h" #include "ggml.h" @@ -13,7 +12,6 @@ #include #include #include -#include #include #include @@ -123,14 +121,10 @@ 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) { @@ -430,8 +424,6 @@ 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, @@ -440,8 +432,6 @@ struct common_sampler * common_sampler_init( /* .prev = */ ring_buffer(std::max(32, params.n_prev)), /* .cur = */ {}, /* .cur_p = */ {}, - /* .speculative_seed = */ speculative_seed, - /* .speculative_rng = */ std::mt19937(speculative_seed), }; return result; @@ -525,8 +515,6 @@ 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, }; } @@ -547,8 +535,6 @@ 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; } @@ -719,79 +705,6 @@ std::vector common_sampler_sample_and_accept_n(struct common_sample return result; } -std::vector common_sampler_sample_and_accept_n( - struct common_sampler * gsmpl, - struct llama_context * ctx, - const std::vector & idxs, - const llama_tokens & draft, - const std::vector & dists, - bool grammar_first) { - GGML_ASSERT(idxs.size() == draft.size() + 1); - GGML_ASSERT(dists.size() == draft.size()); - - std::vector result; - result.reserve(idxs.size()); - - std::uniform_real_distribution 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 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 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 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 common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) { std::vector idxs(draft.size() + 1); for (size_t i = 0; i < idxs.size(); ++i) { @@ -801,19 +714,6 @@ std::vector common_sampler_sample_and_accept_n(struct common_sample return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first); } -std::vector common_sampler_sample_and_accept_n( - struct common_sampler * gsmpl, - struct llama_context * ctx, - const llama_tokens & draft, - const std::vector & dists, - bool grammar_first) { - std::vector 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 index 5f3273eae4b..ced3c8364b3 100644 --- a/common/sampling.h +++ b/common/sampling.h @@ -33,7 +33,6 @@ // struct common_sampler; -struct common_speculative_token_dist; // llama_sampler API overloads @@ -86,23 +85,6 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co // std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first = false); -// maximal-coupling verification for stochastic speculative decoding -std::vector common_sampler_sample_and_accept_n( - struct common_sampler * gsmpl, - struct llama_context * ctx, - const std::vector & idxs, - const llama_tokens & draft, - const std::vector & dists, - bool grammar_first = false); - -// assume idxs == [ 0, 1, 2, ..., draft.size() ] -std::vector common_sampler_sample_and_accept_n( - struct common_sampler * gsmpl, - struct llama_context * ctx, - const llama_tokens & draft, - const std::vector & dists, - bool grammar_first = false); - // assume idxs == [ 0, 1, 2, ..., draft.size() ] std::vector 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 index 113578c7aa2..59bb9c13c59 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -14,12 +14,11 @@ #include #include -#include -#include +#include #include #include #include -#include +#include #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__) @@ -927,8 +926,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { bool is_dflash2 = false; int32_t selector_top_k = 0; - std::vector selector_rng; - std::vector selector_reset; // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout const bool is_dspark; @@ -1007,9 +1004,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { 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 && !is_dflash2) { @@ -1063,8 +1057,6 @@ 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. " @@ -1231,72 +1223,32 @@ 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); + predecessor = (int32_t) std::distance(scores, + std::max_element(scores, scores + selector_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 < 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]; + sum += std::exp(scores[k] - scores[predecessor]); } - for (float & p : dist.probs) { - p /= sum; - } - std::discrete_distribution sample(dist.probs.begin(), dist.probs.end()); - predecessor = sample(selector_rng[seq_id]); - if (dist.probs[predecessor] < params.p_min) { + if (1.0f / sum < params.p_min) { break; } - 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)); - 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 < selector_top_k; ++k) { - sum += std::exp(scores[k] - scores[predecessor]); - } - if (1.0f / sum < params.p_min) { - break; - } - } - 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 (dp.dists) { - dp.dists->clear(); - } } continue; } diff --git a/common/speculative.h b/common/speculative.h index f7dd6526737..12ae31b7de5 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -5,11 +5,6 @@ struct common_speculative; -struct common_speculative_token_dist { - llama_tokens ids; - std::vector probs; -}; - // comma separated list the provided types std::string common_speculative_type_name_str(const std::vector & types); @@ -65,12 +60,6 @@ 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 * 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/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index d074f270f97..487ae03abfa 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -66,8 +66,7 @@ int main(int argc, char ** argv) { llama_context * ctx_dft = params.speculative.draft.ctx_dft; // check if the context supports partial sequence removal - const auto ctx_tgt_seq_rm_type = common_context_can_seq_rm(ctx_tgt); - const bool use_ckpt_tgt = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; + const bool use_ckpt_tgt = common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; const bool use_ckpt_dft = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL; if (use_ckpt_tgt) { @@ -153,7 +152,6 @@ int main(int argc, char ** argv) { llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1); llama_tokens draft; - std::vector dists; common_prompt_checkpoint ckpt; @@ -194,9 +192,6 @@ int main(int argc, char ** argv) { /* .id_last = */ id_last, /* .prompt = */ &prompt_tgt, /* .result = */ &draft, // output - /* .dists = */ &dists, - /* .temperature = */ params.sampling.temp, - /* .seed = */ common_sampler_get_seed(smpl.get()), }; common_speculative_draft(spec); @@ -260,12 +255,7 @@ int main(int argc, char ** argv) { // available logits from the batch and sample the next token until we run out of logits or the sampler // disagrees with the draft // - const bool can_rollback = - ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART || - (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() <= llama_n_rs_seq(ctx_tgt)); - auto ids = can_rollback && params.sampling.temp > 0.0f && dists.size() == draft.size() - ? common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft, dists) - : common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft); + auto ids = common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft); //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str()); @@ -278,7 +268,6 @@ int main(int argc, char ** argv) { LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, n_draft); draft = std::move(ids); - dists.clear(); { ckpt.load_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); @@ -336,7 +325,6 @@ int main(int argc, char ** argv) { // clear the draft since it has been consumed draft.clear(); - dists.clear(); { LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index df11f31b8af..842e4203cd2 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -209,7 +209,6 @@ struct server_slot { common_speculative * spec; llama_tokens spec_draft; - std::vector spec_dists; llama_tokens spec_prompt; std::vector spec_i_batch; common_prompt_checkpoint spec_ckpt; @@ -338,7 +337,6 @@ struct server_slot { if (can_speculate()) { spec_draft.clear(); - spec_dists.clear(); spec_i_batch.clear(); spec_ckpt.clear(); } @@ -2940,9 +2938,6 @@ struct server_context_impl { /* .id_last = */ slot.sampled, /* .prompt = */ &slot.spec_prompt, /* .result = */ &slot.spec_draft, - /* .dists = */ &slot.spec_dists, - /* .temperature = */ slot.task->params.sampling.temp, - /* .seed = */ common_sampler_get_seed(slot.smpl.get()), }; drafting.push_back(&slot); @@ -3822,13 +3817,7 @@ struct server_context_impl { common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); - const bool can_rollback = - ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART || - (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && n_draft <= llama_n_rs_seq(ctx_tgt)); - auto accepted = can_rollback && slot.task->params.sampling.temp > 0.0f && - slot.spec_dists.size() == slot.spec_draft.size() - ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, slot.spec_dists) - : common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); + auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1); @@ -3849,7 +3838,6 @@ struct server_context_impl { // partial acceptance is not supported by the context -> truncate the draft and restore the state slot.spec_is_replay = true; slot.spec_draft = std::move(accepted); - slot.spec_dists.clear(); const auto & ckpt = slot.spec_ckpt; @@ -3877,7 +3865,6 @@ struct server_context_impl { common_speculative_accept(spec.get(), slot.id, accepted.size() - 1); slot.spec_draft = std::move(accepted); - slot.spec_dists.clear(); } const auto ids = std::move(slot.spec_draft); From 2474373ec2e903ec13cac88a6bf22927e65ac6d0 Mon Sep 17 00:00:00 2001 From: SubSir Date: Sat, 22 Aug 2026 15:56:02 -0700 Subject: [PATCH 05/16] Refactor code structure Assisted-by: Claude Opus 5 --- src/llama-graph.h | 2 - src/llama-model.cpp | 1 - src/models/dflash.cpp | 188 +++++++++++++++++++++--------------------- src/models/models.h | 3 - 4 files changed, 96 insertions(+), 98 deletions(-) diff --git a/src/llama-graph.h b/src/llama-graph.h index 605733f32b0..94324c7457e 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1326,8 +1326,6 @@ struct llm_graph_context { void build_sampling() const; - virtual void build_post_sampling() const {} - // // dense (out) // diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e31f38ea413..0d74a2135b6 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2459,7 +2459,6 @@ 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/models/dflash.cpp b/src/models/dflash.cpp index 776d34dd2c6..3f3965cb57a 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -265,7 +265,7 @@ ggml_tensor * llama_model_dflash::graph::build_inp_embd_enc() const { // DFlash Encoder: processes target model features through feature fusion layer template <> -llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model) { +llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { ggml_tensor * cur = build_inp_embd_enc(); cur = build_lora_mm(model.fc, cur, model.fc_s); @@ -446,11 +446,102 @@ static ggml_tensor * build_dflash2_conv( 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(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 candidate_ids(block_size); + std::vector 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 = g.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_pos = 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_pos, 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); + g.cb(packed, "dflash2_lattice", -1); + res->t_h_nextn = packed; + ggml_build_forward_expand(g.gf, packed); +} + // 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::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model) { +llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { const int64_t n_embd_head = hparams.n_embd_head_v(); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); @@ -693,98 +784,11 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra if (model.dspark_markov_w1) { build_dspark_markov_head(*this, model, inp_tokens); } -} - -template -void llama_model_dflash::graph::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(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 candidate_ids(block_size); - std::vector 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); + // DFlash2: pack the selector lattice for the CPU-side walk + if (model.dflash_selector_hidden) { + build_dflash2_selector(*this, model, inp_tokens); } - - 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): diff --git a/src/models/models.h b/src/models/models.h index 969bec894c7..180b30a46d7 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1342,12 +1342,9 @@ struct llama_model_dflash : public llama_model_base { template 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 { From bfaa12563150036bea88853b3974dcd642adbac1 Mon Sep 17 00:00:00 2001 From: SubSir Date: Sat, 22 Aug 2026 16:54:30 -0700 Subject: [PATCH 06/16] Delete embedding scaling Assisted-by: Claude Opus 5 --- src/models/dflash.cpp | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 3f3965cb57a..933da7fd85d 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -10,7 +10,6 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { 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); @@ -646,9 +645,6 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra 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)); From cf095c83421a636f609d36637a2c59f4035c8132 Mon Sep 17 00:00:00 2001 From: SubSir Date: Sat, 22 Aug 2026 17:05:07 -0700 Subject: [PATCH 07/16] Gate output transforms on DFlash2 Assisted-by: Claude Opus 5 --- src/models/dflash.cpp | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 933da7fd85d..57a05ccb0d9 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -747,13 +747,18 @@ llama_model_dflash::graph::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); + // DFlash2 feeds these logits to the selector, so they carry the target's + // output transforms. DFlash1 and DSpark read them through the sampler and + // are left untouched. + if (model.dflash_selector_hidden) { + 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 From 7ea40ee98acb416787863aee935dbb99491acad5 Mon Sep 17 00:00:00 2001 From: SubSir Date: Sat, 22 Aug 2026 21:45:34 -0700 Subject: [PATCH 08/16] Optimize Dflash 2 cost Assisted-by: Claude Opus 5 --- ggml/src/ggml-cuda/top-k.cu | 117 ++++++++++++++++++++++++++++ src/models/dflash.cpp | 150 ++++++++++++++++++++---------------- tests/test-backend-ops.cpp | 19 ++++- 3 files changed, 220 insertions(+), 66 deletions(-) diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 9681cd29333..6e839d056c8 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -48,6 +48,118 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE + +// Two-stage top-k, for rows too wide to select one at a time. +// +// Tiling first cannot drop a winner: a global top-k element has at most k-1 +// larger elements anywhere, so at most k-1 inside its own tile. What survives +// fits the existing bitonic argsort, which finishes every row in one launch. +// +// Keys pack the value and the column into 64 bits, so a max reduction resolves +// the ordering and the tie-break at once and every key in a tile is distinct - +// which is what makes clearing the entry equal to the running maximum remove +// exactly one element. + +// 1024 is argsort_f32_i32_cuda_bitonic's row limit. It bounds what the second +// stage can be handed, and it is also where this path starts to pay: below it +// the bitonic sort already takes the row whole and is faster (measured on +// H200 and A10G, k=16, 16 rows: 5 us vs 18 us at 16 columns, crossing over by +// 1000). +#define TOPK_CAND 1024 + +// Block size measured best at 256 on both cards and every width; 512 and 1024 +// lose more to occupancy than they gain in reduction depth. The tile size does +// not have one winner - narrow rows prefer 4096, and both cards agree on 8192 +// from 65536 columns up (A10G, 200000x16: 132 us at 4096 vs 120 us at 8192). +#define TOPK_BLOCK 256 +#define TOPK_TILE_WIDE 8192 +#define TOPK_TILE 4096 + +template +static __global__ void topk_tile(const float * src, float * cand_val, int * cand_idx, + const int ncols, const int ntiles, const int k) { + __shared__ uint64_t smem[BLOCK]; + + const int row = blockIdx.x / ntiles; + const int tile = blockIdx.x % ntiles; + const float * row_ptr = src + (size_t) row * ncols; + + uint64_t keys[TILE / BLOCK]; +#pragma unroll + for (int i = 0; i < TILE / BLOCK; ++i) { + const int col = tile * TILE + threadIdx.x + i * BLOCK; + uint32_t b = col < ncols ? __float_as_uint(row_ptr[col]) : 0; + b = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + keys[i] = col < ncols ? (((uint64_t) b << 32) | (uint32_t) (ncols - 1 - col)) : 0; + } + + const size_t out = ((size_t) row * ntiles + tile) * k; + for (int j = 0; j < k; ++j) { + uint64_t local = 0; +#pragma unroll + for (int i = 0; i < TILE / BLOCK; ++i) { + local = max(local, keys[i]); + } + smem[threadIdx.x] = local; + __syncthreads(); + for (int s = BLOCK / 2; s > 0; s >>= 1) { + if (threadIdx.x < s) { + smem[threadIdx.x] = max(smem[threadIdx.x], smem[threadIdx.x + s]); + } + __syncthreads(); + } + const uint64_t best = smem[0]; + if (threadIdx.x == 0) { + const int col = ncols - 1 - (int) (best & 0xFFFFFFFFu); + cand_val[out + j] = best ? row_ptr[col] : -INFINITY; + cand_idx[out + j] = best ? col : 0; + } +#pragma unroll + for (int i = 0; i < TILE / BLOCK; ++i) { + if (keys[i] == best) { + keys[i] = 0; + } + } + __syncthreads(); + } +} + +// The argsort ranks candidates; turn its positions back into columns. +static __global__ void topk_unmap(const int * cand_idx, const int * order, int * dst, + const int ncand, const int k) { + for (int i = threadIdx.x; i < k; i += blockDim.x) { + dst[(size_t) blockIdx.x * k + i] = cand_idx[(size_t) blockIdx.x * ncand + order[(size_t) blockIdx.x * ncand + i]]; + } +} + +static bool ggml_cuda_top_k_tiled(ggml_cuda_pool & pool, const float * src, int * dst, + const int ncols, const int nrows, const int k, + cudaStream_t stream) { + // Narrow rows are already handled whole by the bitonic sort below. + const int tile = ncols >= 65536 ? TOPK_TILE_WIDE : TOPK_TILE; + const int ntiles = (ncols + tile - 1) / tile; + const int ncand = ntiles * k; + if (ncols <= TOPK_CAND || ncand > TOPK_CAND) { + return false; + } + + ggml_cuda_pool_alloc cand_val(pool, (size_t) nrows * ncand); + ggml_cuda_pool_alloc cand_idx(pool, (size_t) nrows * ncand); + ggml_cuda_pool_alloc order (pool, (size_t) nrows * ncand); + + if (tile == TOPK_TILE_WIDE) { + topk_tile<<>>( + src, cand_val.get(), cand_idx.get(), ncols, ntiles, k); + } else { + topk_tile<<>>( + src, cand_val.get(), cand_idx.get(), ncols, ntiles, k); + } + argsort_f32_i32_cuda_bitonic(cand_val.get(), order.get(), ncand, nrows, + GGML_SORT_ORDER_DESC, stream); + topk_unmap<<>>(cand_idx.get(), order.get(), dst, ncand, k); + return true; +} + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -63,6 +175,11 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t nrows = ggml_nrows(src0); const int64_t k = dst->ne[0]; ggml_cuda_pool & pool = ctx.pool(); + + if (ggml_cuda_top_k_tiled(pool, src0_d, dst_d, ncols, nrows, k, stream)) { + return; + } + #ifdef CUB_TOP_K_AVAILABLE // TODO: Switch to `DeviceSegmentedTopK` for multi-row TopK once implemented // https://github.com/NVIDIA/cccl/issues/6391 diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 57a05ccb0d9..a0b801a586c 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -406,14 +406,34 @@ static ggml_tensor * build_dflash2_conv( 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); + // Both arrive contiguous from the caller in the common case; ggml_cont + // would still copy, and this runs 4x per layer. + 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 * 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]); + // All taps' coefficients live in one tensor already, so the contiguous copy + // and the group broadcast are done once instead of per tap. The (coeff + + // base) form is kept fused: distributing it over the multiply would change + // the rounding. + 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); + ggml_tensor * result = nullptr; for (int64_t tap = 0; tap < kernel_size; ++tap) { ggml_tensor * values = blocks; @@ -430,15 +450,11 @@ static ggml_tensor * build_dflash2_conv( } 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 * 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 * 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; } @@ -466,68 +482,72 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo const int64_t tokens_per_block = n_tokens / n_blocks; const int64_t block_size = std::min(tokens_per_block, hparams.dflash_block_size); - ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k); + 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_tensor * unary = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens); + ggml_tensor * hidden = 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 * cand3 = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks); + ggml_tensor * unary3 = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks); + ggml_tensor * hid3 = ggml_reshape_3d(ctx0, hidden, rank, tokens_per_block, n_blocks); + + // Score a run of block positions at once. A position's score depends only + // on the candidate sets at pos-1 and pos, both of which top-k already + // produced, so the positions carry no dependency on one another and the + // whole run is one batched matmul instead of one per position. + auto score_run = [&](int64_t beg_pos, int64_t n_pos, ggml_tensor * pred_ids) { + ggml_tensor * ids = ggml_cont(ctx0, ggml_view_3d(ctx0, cand3, top_k, n_pos, n_blocks, + cand3->nb[1], cand3->nb[2], beg_pos * cand3->nb[1])); + ggml_tensor * un = ggml_cont(ctx0, ggml_view_3d(ctx0, unary3, top_k, n_pos, n_blocks, + unary3->nb[1], unary3->nb[2], beg_pos * unary3->nb[1])); + ggml_tensor * hid = ggml_cont(ctx0, ggml_view_3d(ctx0, hid3, rank, n_pos, n_blocks, + hid3->nb[1], hid3->nb[2], beg_pos * hid3->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, ids, 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 * hid4 = ggml_reshape_4d(ctx0, hid, rank, 1, n_pos, n_blocks); + ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, hid4, 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 * un4 = ggml_reshape_4d(ctx0, un, top_k, 1, n_pos, n_blocks); + score = ggml_add(ctx0, score, ggml_repeat(ctx0, un4, score)); - std::vector candidate_ids(block_size); - std::vector 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 = g.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 * row = ggml_concat(ctx0, + ggml_cast(ctx0, ids, 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); - for (int64_t pos = 1; pos < block_size; ++pos) { - ggml_tensor * ids = candidate_ids[pos]; - ggml_tensor * unary_pos = 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_pos, 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); + 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, cand3, top_k, block_size - 2, n_blocks, + cand3->nb[1], cand3->nb[2], cand3->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); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 3349a64b174..9d9e135a485 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9637,6 +9637,17 @@ static std::vector> make_test_cases_eval() { } } } + for (int k : {4, 8, 16, 32}) { + for (int nrows : {1, 8, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {151936, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8193, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, nrows, 1, 1}, k, true)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k, true)); + } + } + for (int k : {1, 2, 3, 7, 15}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {16, 10, 10, 10}, k)); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {60, 10, 10, 10}, k)); @@ -10321,7 +10332,13 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {200000, 16, 1, 1})); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2, 1, 1, 1}, 1)); - for (auto k : {1, 10, 40, 400}) { + // Crossover sweep: where does tiling start paying for itself? + for (auto cols : {4096, 8192, 12288, 16384, 24576, 32768, 65536, 131072}) { + for (auto nrows : {1, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 16)); + } + } + for (auto k : {1, 4, 8, 10, 16, 32, 40, 400}) { for (auto nrows : {1, 16}) { for (auto cols : {k, 1000, 65000, 200000}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k)); From c184d9a7a113f7d6e9a5a72948751d35a37c980b Mon Sep 17 00:00:00 2001 From: SubSir Date: Sun, 23 Aug 2026 15:58:03 -0700 Subject: [PATCH 09/16] Avoid using atoi Assisted-by: Claude Opus 5 --- common/speculative.cpp | 7 +++---- include/llama.h | 1 + src/llama-model.cpp | 4 ++++ 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 59bb9c13c59..b9d0c095bf8 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -970,11 +970,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; - } } + + selector_top_k = llama_model_dflash_selector_top_k(model_dft); + is_dflash2 = selector_top_k > 0; mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); diff --git a/include/llama.h b/include/llama.h index 177fc10a913..22615a95031 100644 --- a/include/llama.h +++ b/include/llama.h @@ -582,6 +582,7 @@ extern "C" { LLAMA_API int32_t llama_model_n_embd_out (const struct llama_model * model); LLAMA_API int32_t llama_model_n_layer (const struct llama_model * model); LLAMA_API int32_t llama_model_n_layer_nextn(const struct llama_model * model); + LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model); LLAMA_API int32_t llama_model_n_head (const struct llama_model * model); LLAMA_API int32_t llama_model_n_head_kv (const struct llama_model * model); LLAMA_API int32_t llama_model_n_swa (const struct llama_model * model); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0d74a2135b6..4115f1fad7d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2535,6 +2535,10 @@ int32_t llama_model_n_layer_nextn(const llama_model * model) { return model->hparams.n_layer_nextn; } +int32_t llama_model_dflash_selector_top_k(const llama_model * model) { + return model->hparams.dflash_selector_top_k; +} + int32_t llama_model_n_head(const llama_model * model) { return model->hparams.n_head(); } From d1a522fc89c96d1a3057e35681f0c4859810623c Mon Sep 17 00:00:00 2001 From: SubSir Date: Sun, 23 Aug 2026 15:58:03 -0700 Subject: [PATCH 10/16] Modify comments Assisted-by: Claude Opus 5 --- ggml/src/ggml-cuda/top-k.cu | 27 +++++---------------------- src/llama-context.cpp | 2 ++ src/models/dflash.cpp | 19 +++++-------------- tests/test-backend-ops.cpp | 2 +- 4 files changed, 13 insertions(+), 37 deletions(-) diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 6e839d056c8..a75c6d540e0 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -49,28 +49,11 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE -// Two-stage top-k, for rows too wide to select one at a time. -// -// Tiling first cannot drop a winner: a global top-k element has at most k-1 -// larger elements anywhere, so at most k-1 inside its own tile. What survives -// fits the existing bitonic argsort, which finishes every row in one launch. -// -// Keys pack the value and the column into 64 bits, so a max reduction resolves -// the ordering and the tie-break at once and every key in a tile is distinct - -// which is what makes clearing the entry equal to the running maximum remove -// exactly one element. - -// 1024 is argsort_f32_i32_cuda_bitonic's row limit. It bounds what the second -// stage can be handed, and it is also where this path starts to pay: below it -// the bitonic sort already takes the row whole and is faster (measured on -// H200 and A10G, k=16, 16 rows: 5 us vs 18 us at 16 columns, crossing over by -// 1000). -#define TOPK_CAND 1024 - -// Block size measured best at 256 on both cards and every width; 512 and 1024 -// lose more to occupancy than they gain in reduction depth. The tile size does -// not have one winner - narrow rows prefer 4096, and both cards agree on 8192 -// from 65536 columns up (A10G, 200000x16: 132 us at 4096 vs 120 us at 8192). +// Two-stage top-k for wide rows: a global top-k element has at most k-1 larger +// elements, so at most k-1 inside its own tile and tiling cannot drop a winner. +#define TOPK_CAND 1024 // argsort_f32_i32_cuda_bitonic's row limit + +// measured on H200 and A10G #define TOPK_BLOCK 256 #define TOPK_TILE_WIDE 8192 #define TOPK_TILE 4096 diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 51148b2adcc..b5b4f5dbdaf 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2315,6 +2315,8 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { } if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) { + // the selector runs on block positions only: at most one block per sequence, + // and never more rows than a ubatch holds - so take the tighter bound const uint32_t selector_tokens = std::min( n_tokens, model.hparams.dflash_block_size * cparams.n_seq_max); res += 32*selector_tokens; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index a0b801a586c..bef6076f1fc 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -406,8 +406,7 @@ static ggml_tensor * build_dflash2_conv( const int64_t block_size = n_tokens / n_blocks; ggml_context * ctx0 = g.ctx0; - // Both arrive contiguous from the caller in the common case; ggml_cont - // would still copy, and this runs 4x per layer. + // 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); } @@ -420,10 +419,6 @@ static ggml_tensor * build_dflash2_conv( 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]); - // All taps' coefficients live in one tensor already, so the contiguous copy - // and the group broadcast are done once instead of per tap. The (coeff + - // base) form is kept fused: distributing it over the multiply would change - // the rounding. 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); @@ -496,10 +491,8 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo ggml_tensor * unary3 = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks); ggml_tensor * hid3 = ggml_reshape_3d(ctx0, hidden, rank, tokens_per_block, n_blocks); - // Score a run of block positions at once. A position's score depends only - // on the candidate sets at pos-1 and pos, both of which top-k already - // produced, so the positions carry no dependency on one another and the - // whole run is one batched matmul instead of one per position. + // 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 * ids = ggml_cont(ctx0, ggml_view_3d(ctx0, cand3, top_k, n_pos, n_blocks, cand3->nb[1], cand3->nb[2], beg_pos * cand3->nb[1])); @@ -767,9 +760,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cur = build_lora_mm(output, cur, output_s); - // DFlash2 feeds these logits to the selector, so they carry the target's - // output transforms. DFlash1 and DSpark read them through the sampler and - // are left untouched. + // DFlash2 feeds these logits to the selector, so they need the target's output + // transforms; DFlash1 and DSpark read them through the sampler instead if (model.dflash_selector_hidden) { if (hparams.f_logit_scale != 0.0f) { cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); @@ -806,7 +798,6 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra build_dspark_markov_head(*this, model, inp_tokens); } - // DFlash2: pack the selector lattice for the CPU-side walk if (model.dflash_selector_hidden) { build_dflash2_selector(*this, model, inp_tokens); } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 9d9e135a485..0a8ff97c2aa 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10332,7 +10332,7 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {200000, 16, 1, 1})); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2, 1, 1, 1}, 1)); - // Crossover sweep: where does tiling start paying for itself? + // widths around the tiling threshold for (auto cols : {4096, 8192, 12288, 16384, 24576, 32768, 65536, 131072}) { for (auto nrows : {1, 16}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 16)); From e91bca536c438a29528fe0cc1f12dc2ced61538b Mon Sep 17 00:00:00 2001 From: SubSir Date: Mon, 24 Aug 2026 01:47:37 -0700 Subject: [PATCH 11/16] Move llama_model_dflash_selector_top_k to llama-ext.h Assisted-by: Claude Opus 5 --- include/llama.h | 1 - src/llama-ext.h | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/include/llama.h b/include/llama.h index 22615a95031..177fc10a913 100644 --- a/include/llama.h +++ b/include/llama.h @@ -582,7 +582,6 @@ extern "C" { LLAMA_API int32_t llama_model_n_embd_out (const struct llama_model * model); LLAMA_API int32_t llama_model_n_layer (const struct llama_model * model); LLAMA_API int32_t llama_model_n_layer_nextn(const struct llama_model * model); - LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model); LLAMA_API int32_t llama_model_n_head (const struct llama_model * model); LLAMA_API int32_t llama_model_n_head_kv (const struct llama_model * model); LLAMA_API int32_t llama_model_n_swa (const struct llama_model * model); diff --git a/src/llama-ext.h b/src/llama-ext.h index 35d6e58adfa..92a759b7a0a 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -120,6 +120,8 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx); // model/context data extraction // +LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model); + // returns pointer to the target-model layer indices LLAMA_API const int32_t * llama_model_target_layer_ids (const struct llama_model * model); // returns the number of extracted layers from target model From 64f765f5adefa4620dddda436ce56f1430435536 Mon Sep 17 00:00:00 2001 From: SubSir Date: Mon, 24 Aug 2026 01:47:37 -0700 Subject: [PATCH 12/16] Formatting Assisted-by: Claude Opus 5 --- gguf-py/gguf/constants.py | 14 +++++++------- src/llama-arch.cpp | 11 ++++++----- src/models/dflash.cpp | 6 +++--- 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index e3dc7a7a892..4c556eda61d 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1886,13 +1886,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", - MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base", - MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj", - MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base", - MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj", - MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor", - MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor", - MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden", + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base", + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj", + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base", + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj", + MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor", + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor", + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden", MODEL_TENSOR.D2T: "d2t", } diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 8fd4c8974fb..ec877d4201f 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -334,14 +334,15 @@ static const std::map 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" }, + { 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_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, // sentence-transformers dense modules feature dims { LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" }, diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index bef6076f1fc..f072f3e3345 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -7,11 +7,11 @@ 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); + 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_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); - ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, 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); From f5a7ec15da6add890a5624c0990714498df837a4 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Mon, 24 Aug 2026 11:33:36 -0700 Subject: [PATCH 13/16] Apply patch to fix the mrope bug Assisted-by: Claude Opus 5 --- common/speculative.cpp | 31 +++++++++++++++++++++++++++++-- conversion/qwen.py | 28 ++++++++++++++++++++++++++++ src/llama-model.cpp | 4 ++++ src/models/dflash.cpp | 35 ++++++++++++++++++++--------------- 4 files changed, 81 insertions(+), 17 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index b9d0c095bf8..b0ad48eb483 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -925,6 +925,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { llama_token mask_token_id = 0; bool is_dflash2 = false; + bool is_mrope = false; int32_t selector_top_k = 0; // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout @@ -994,6 +995,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq); + // embd batches on an M-RoPE draft need 4 position rows per token + is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE; + if (is_mrope) { + free(batch_inject.pos); + batch_inject.pos = (llama_pos *) malloc(sizeof(llama_pos) * 4 * llama_n_batch(ctx_dft)); + } + smpls.resize(n_seq); for (auto & s : smpls) { common_params_sampling sparams; @@ -1126,11 +1134,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } // fuse extracted features through DFlash encoder + // M-RoPE drafts read 4 position rows per token from embd batches, so pass them explicitly + std::vector enc_pos; + if (is_mrope) { + enc_pos.resize((size_t) 4 * n_chunk); + for (int32_t i = 0; i < n_chunk; ++i) { + const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + enc_pos[0 * n_chunk + i] = p; + enc_pos[1 * n_chunk + i] = p; + enc_pos[2 * n_chunk + i] = p; + enc_pos[3 * n_chunk + i] = 0; + } + } + llama_batch enc_batch = { /*.n_tokens =*/ n_chunk, /*.token =*/ nullptr, /*.embd =*/ features_buf.data(), - /*.pos =*/ nullptr, + /*.pos =*/ is_mrope ? enc_pos.data() : nullptr, /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, @@ -1151,7 +1172,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { 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) { - batch_inject.pos[i] = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + batch_inject.pos[i] = p; + if (is_mrope) { + batch_inject.pos[1 * n_chunk + i] = p; + batch_inject.pos[2 * n_chunk + i] = p; + batch_inject.pos[3 * n_chunk + i] = 0; + } batch_inject.n_seq_id[i] = 1; batch_inject.seq_id[i][0] = seq_id; batch_inject.logits[i] = false; diff --git a/conversion/qwen.py b/conversion/qwen.py index ab2a0fb4dce..f3944709076 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -717,6 +717,21 @@ def set_gguf_parameters(self): self.gguf_writer.add_sliding_window(sliding_window) self.gguf_writer.add_sliding_window_pattern(is_swa) + # M-RoPE target: the draft ropes on the temporal dim only, so write + # degenerate sections [n_rot/2, 0, 0, 0] + if self._target_uses_mrope(): + head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0]) + + def _target_uses_mrope(self) -> bool: + if self.target_model_dir is None: + return False + with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: + cfg = json.load(f) + cfg = cfg.get("text_config", cfg) + rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {} + return "mrope_section" in rope + @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item @@ -724,10 +739,23 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca name = "model." + name return super().filter_tensors((name, gen)) + _ROPE_PERMUTE_SUFFIXES = ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + ) + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True): return + # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd + if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES): + head_dim = self.hparams["head_dim"] + shape = data_torch.shape + data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape) + if name in ( "model.candidate_selector.predecessor_codebook", "model.candidate_selector.successor_codebook", diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4115f1fad7d..de3b608f6ba 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2735,6 +2735,10 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_DFLASH: + // drafts for M-RoPE targets carry rope sections and follow the target's temporal dim + if (const auto & s = model->hparams.rope_sections; s[0] || s[1] || s[2] || s[3]) { + return LLAMA_ROPE_TYPE_MROPE; + } // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index f072f3e3345..68b8c1f1d71 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -11,6 +11,9 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { hparams.f_final_logit_softcapping = 0.0f; ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + // drafts for M-RoPE targets carry degenerate sections [n_rot/2, 0, 0, 0] + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, 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); @@ -573,6 +576,20 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra const float kq_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, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + }; + // KV cache injection if (ubatch.embd) { auto inp = std::make_unique(n_embd); @@ -595,11 +612,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra 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); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + Kcur = build_rope(Kcur, inp_pos); cb(Kcur, "Kcur_injected", il); cb(Vcur, "Vcur_injected", il); @@ -686,16 +699,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + Qcur = build_rope(Qcur, inp_pos); + Kcur = build_rope(Kcur, inp_pos); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); From f7aadef0932e47d66a4349245957e81126a7c734 Mon Sep 17 00:00:00 2001 From: SubSir Date: Mon, 24 Aug 2026 16:00:05 -0700 Subject: [PATCH 14/16] fix ci Assisted-by: Claude Opus 5 --- src/models/dflash.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 68b8c1f1d71..cd80181eb2e 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -417,7 +417,6 @@ static ggml_tensor * build_dflash2_conv( 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]); From 11f45ed3439894c93a34e8820341aff0d5cc3c2d Mon Sep 17 00:00:00 2001 From: SubSir Date: Wed, 26 Aug 2026 13:38:21 -0700 Subject: [PATCH 15/16] Fix graph number calculation Assisted-by: Claude Opus 5 --- src/llama-context.cpp | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index b5b4f5dbdaf..ac2be97a653 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2307,6 +2307,10 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_MINIMAX_01 || model.arch == LLM_ARCH_MINIMAX_M3) { res = std::max(n_tokens * 40, 32u * model.n_tensors()); + } else if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) { + // DFlash2's convolutions and selector are shape work rather than matmuls, + // so they cost ~8.6 nodes per tensor against ~5.9 for a plain DFlash draft + res = std::max(1024u, 12u*model.n_tensors()); } else { res = std::max(1024u, 8u*model.n_tensors()); for (const auto & lora : model.loras) { @@ -2314,14 +2318,6 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { } } - if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) { - // the selector runs on block positions only: at most one block per sequence, - // and never more rows than a ubatch holds - so take the tighter bound - const uint32_t selector_tokens = std::min( - 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) { From 2f3923bc81346046aa5765dda15fb28497f49ac6 Mon Sep 17 00:00:00 2001 From: SubSir Date: Wed, 26 Aug 2026 13:53:51 -0700 Subject: [PATCH 16/16] rename hid and unary Assisted-by: Claude Opus 5 --- src/models/dflash.cpp | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index cd80181eb2e..2ca0491a37d 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -485,44 +485,44 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo 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 * hidden = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd); + 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 * cand3 = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks); - ggml_tensor * unary3 = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks); - ggml_tensor * hid3 = ggml_reshape_3d(ctx0, hidden, rank, tokens_per_block, n_blocks); + 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 * ids = ggml_cont(ctx0, ggml_view_3d(ctx0, cand3, top_k, n_pos, n_blocks, - cand3->nb[1], cand3->nb[2], beg_pos * cand3->nb[1])); - ggml_tensor * un = ggml_cont(ctx0, ggml_view_3d(ctx0, unary3, top_k, n_pos, n_blocks, - unary3->nb[1], unary3->nb[2], beg_pos * unary3->nb[1])); - ggml_tensor * hid = ggml_cont(ctx0, ggml_view_3d(ctx0, hid3, rank, n_pos, n_blocks, - hid3->nb[1], hid3->nb[2], beg_pos * hid3->nb[1])); + 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, ids, top_k * n_pos * n_blocks)), + 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 * hid4 = ggml_reshape_4d(ctx0, hid, rank, 1, n_pos, n_blocks); - ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, hid4, predecessor)); + 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 * un4 = ggml_reshape_4d(ctx0, un, top_k, 1, n_pos, n_blocks); - score = ggml_add(ctx0, score, ggml_repeat(ctx0, un4, score)); + 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, ids, GGML_TYPE_F32), + 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); }; @@ -539,8 +539,8 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo } if (block_size > 2) { ggml_tensor * prev_ids = ggml_reshape_1d(ctx0, - ggml_cont(ctx0, ggml_view_3d(ctx0, cand3, top_k, block_size - 2, n_blocks, - cand3->nb[1], cand3->nb[2], cand3->nb[1])), + 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); }