From 7ebbb906d2ca9a3a5a9590819672f0b205c56a18 Mon Sep 17 00:00:00 2001 From: Samuel Oliveira Alves <107287165+SamuelOliveirads@users.noreply.github.com> Date: Mon, 10 Aug 2026 03:46:03 -0300 Subject: [PATCH] Initial implementation of DSpark (#2280) * Implement initial arch for DSpark * feat: Add Dspark architecture support * avoid to many splits in graph and improve rope logic --- common/common.cpp | 32 ++++- common/common.h | 5 + common/speculative-dflash-impl.h | 29 ++++- common/speculative.cpp | 47 +++++-- convert_hf_to_gguf.py | 60 +++++++++ examples/main/main.cpp | 3 +- examples/server/server-context.cpp | 3 +- gguf-py/gguf/constants.py | 10 ++ src/graphs/build_deepseek4.cpp | 195 ++++++++++++++++++++++++++++- src/graphs/build_dflash.cpp | 112 ++++++++++++++++- src/llama-arch.cpp | 6 + src/llama-arch.h | 6 + src/llama-build-context.cpp | 1 + src/llama-build-context.h | 6 + src/llama-context.h | 4 +- src/llama-dflash.cpp | 46 +++++-- src/llama-dsv4.cpp | 23 +++- src/llama-hparams.cpp | 45 ++++++- src/llama-hparams.h | 2 + src/llama-load-tensors.cpp | 118 +++++++++++++++++ src/llama-model.cpp | 43 +++++++ src/llama-model.h | 4 + src/llama-spec-features-dflash.cpp | 16 ++- src/llama-spec-features-dflash.h | 2 + src/llama.cpp | 14 ++- 25 files changed, 776 insertions(+), 56 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index bb1d7d8bb..f241371fb 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -96,6 +96,14 @@ common_time_meas::~common_time_meas() { } } +bool common_speculative_type_is_dflash_family(enum common_speculative_type type) { + return type == COMMON_SPECULATIVE_TYPE_DFLASH || type == COMMON_SPECULATIVE_TYPE_DSPARK; +} + +bool common_speculative_type_uses_target_features(enum common_speculative_type type) { + return type == COMMON_SPECULATIVE_TYPE_MTP || common_speculative_type_is_dflash_family(type); +} + bool common_speculative_type_is_self_spec(enum common_speculative_type type) { switch (type) { case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: @@ -202,6 +210,20 @@ bool common_params_speculative::has_stage_type(common_speculative_type stage_typ }); } +bool common_params_speculative::has_dflash_family_stage() const { + const auto resolved = get_resolved_stages(); + return std::any_of(resolved.begin(), resolved.end(), [](const common_speculative_stage_params & stage) { + return common_speculative_type_is_dflash_family(stage.type); + }); +} + +bool common_params_speculative::uses_target_features() const { + const auto resolved = get_resolved_stages(); + return std::any_of(resolved.begin(), resolved.end(), [](const common_speculative_stage_params & stage) { + return common_speculative_type_uses_target_features(stage.type); + }); +} + void common_params_speculative::remove_stage_type(common_speculative_type stage_type) { stages.erase(std::remove_if(stages.begin(), stages.end(), [stage_type](const common_speculative_stage_params & stage) { return stage.type == stage_type; @@ -219,7 +241,7 @@ bool common_params_speculative::has_composite_stage_chain() const { bool common_params_speculative::needs_dft_model() const { return has_stage_type(COMMON_SPECULATIVE_TYPE_DRAFT) || - has_stage_type(COMMON_SPECULATIVE_TYPE_DFLASH) || + has_dflash_family_stage() || (has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) && has_dft()); } @@ -295,12 +317,12 @@ bool common_speculative_validate_chain(const common_params_speculative & params, return fail("speculative stage has n_min greater than n_max"); } - if ((stage.type == COMMON_SPECULATIVE_TYPE_DRAFT || stage.type == COMMON_SPECULATIVE_TYPE_DFLASH) && !params.has_dft()) { + if ((stage.type == COMMON_SPECULATIVE_TYPE_DRAFT || common_speculative_type_is_dflash_family(stage.type)) && !params.has_dft()) { return fail(common_speculative_type_to_str(stage.type) + " speculative stage requires a draft model or draft params"); } - if (stage.type == COMMON_SPECULATIVE_TYPE_DFLASH && stage_params.dflash_cross_ctx < 1) { - return fail("dflash speculative stage requires cross_ctx >= 1"); + if (common_speculative_type_is_dflash_family(stage.type) && stage_params.dflash_cross_ctx < 1) { + return fail(common_speculative_type_to_str(stage.type) + " speculative stage requires cross_ctx >= 1"); } } @@ -3366,7 +3388,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param " cpu serialise architecture state via host storage; re-decode on rejection\n" " --recurrent-ckpt-mode remains as a deprecated alias" }); options.push_back({ "*", "--spec-type SPEC[:k=v,...]", "canonical speculative stage entry; repeat for a supported two-stage chain.\n" - "types: none, draft, dflash, mtp, ngram-cache, ngram-simple, ngram-map-k, ngram-map-k4v, ngram-mod, suffix\n" + "types: none, draft, dflash, dspark, mtp, ngram-cache, ngram-simple, ngram-map-k, ngram-map-k4v, ngram-mod, suffix\n" "canonical keys: n_max,n_min,p_min,heads,cross_ctx,ngram_size_n,ngram_size_m,ngram_min_hits,suffix_min_match_len,suffix_max_depth,suffix_corpus\n" "MTP heads: heads=1 is the default; heads>1 and heads=0 (all model heads) are experimental\n" "for comma-bearing string values, quote the value inside the stage payload for normal shell use\n" diff --git a/common/common.h b/common/common.h index 51b00497b..086d98d67 100644 --- a/common/common.h +++ b/common/common.h @@ -150,6 +150,7 @@ enum common_speculative_type { COMMON_SPECULATIVE_TYPE_NONE, // no speculative decoding COMMON_SPECULATIVE_TYPE_DRAFT, // draft model COMMON_SPECULATIVE_TYPE_DFLASH, // DFlash draft model + COMMON_SPECULATIVE_TYPE_DSPARK, COMMON_SPECULATIVE_TYPE_MTP, // MTP model COMMON_SPECULATIVE_TYPE_EAGLE3, // eagle draft model COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding @@ -164,6 +165,8 @@ enum common_speculative_type { std::string common_speculative_type_name_str(); enum common_speculative_type common_speculative_type_from_name(const std::string & name); std::string common_speculative_type_to_str(enum common_speculative_type type); +bool common_speculative_type_is_dflash_family(enum common_speculative_type type); +bool common_speculative_type_uses_target_features(enum common_speculative_type type); bool common_speculative_type_is_self_spec(enum common_speculative_type type); struct common_speculative_stage_params { @@ -268,6 +271,8 @@ struct common_params_speculative { common_params_speculative with_stage_overrides(const common_speculative_stage_params & stage) const; bool has_stage_chain() const; bool has_stage_type(common_speculative_type stage_type) const; + bool has_dflash_family_stage() const; + bool uses_target_features() const; void remove_stage_type(common_speculative_type stage_type); bool has_composite_stage_chain() const; bool needs_dft_model() const; diff --git a/common/speculative-dflash-impl.h b/common/speculative-dflash-impl.h index 746c8130b..f957fe3f0 100644 --- a/common/speculative-dflash-impl.h +++ b/common/speculative-dflash-impl.h @@ -8,6 +8,12 @@ static bool common_speculative_are_dflash_compatible( const llama_model * model_tgt, const llama_model * model_dft) { + const char * draft_arch = model_dft != nullptr ? llama_model_arch_string(model_dft) : nullptr; + if (model_tgt == nullptr || model_dft == nullptr || draft_arch == nullptr || + (std::strcmp(draft_arch, "dflash") != 0 && std::strcmp(draft_arch, "dflash-draft") != 0)) { + return false; + } + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); @@ -82,6 +88,7 @@ struct common_speculative_state_dflash : public common_speculative_state { int32_t mask_token_id = -1; int32_t n_target_features = 0; int32_t cross_ctx = 0; + bool is_dspark = false; bool ready = false; std::vector target_layer_ids; @@ -114,6 +121,15 @@ struct common_speculative_state_dflash : public common_speculative_state { const llama_model * model_tgt = llama_get_model(ctx_tgt); const llama_model * model_dft = llama_get_model(ctx_dft); + is_dspark = type == COMMON_SPECULATIVE_TYPE_DSPARK; + const bool has_dspark_head = llama_model_dflash_has_dspark_head(model_dft); + if (is_dspark != has_dspark_head) { + LOG_ERR("%s: %s stage requires %s DSpark Markov tensors\n", __func__, + is_dspark ? "dspark" : "dflash", + is_dspark ? "complete" : "no"); + return; + } + if (!common_speculative_are_dflash_compatible(model_tgt, model_dft)) { LOG_ERR("%s: DFlash draft model vocab/tokenizer is incompatible with the target model\n", __func__); return; @@ -214,6 +230,7 @@ struct common_speculative_state_dflash : public common_speculative_state { ready = true; llama_set_dflash_visible_cross_ctx(ctx_dft, this->cross_ctx); + llama_set_dflash_dspark(ctx_dft, is_dspark); LOG_INF("%s: DFlash context ready (n_ctx=%d, block_size=%d, cross_ctx=%d, n_target_features=%d, n_target_layers=%d)\n", __func__, llama_n_ctx(ctx_dft), block_size, this->cross_ctx, n_target_features, n_target_layers); } @@ -246,7 +263,8 @@ struct common_speculative_state_dflash : public common_speculative_state { return; } - const int32_t n_keep = std::min(params.n_max, block_size - 1); + const int32_t max_draft_tokens = is_dspark ? block_size : block_size - 1; + const int32_t n_keep = std::min(params.n_max, max_draft_tokens); if (n_keep <= 0) { return; } @@ -280,12 +298,12 @@ struct common_speculative_state_dflash : public common_speculative_state { llama_kv_cache_clear(ctx_dft); batch.n_tokens = 0; - const int32_t batch_len = n_keep + 1; + const int32_t batch_len = is_dspark ? n_keep : n_keep + 1; const llama_pos draft_pos_base = last_target_pos >= 0 ? last_target_pos + 1 : (llama_pos) target_window_rows; const llama_pos seed_pos = last_target_pos >= 0 ? last_target_pos : draft_pos_base - 1; - common_batch_add(batch, id_last, seed_pos, { 0 }, false); + common_batch_add(batch, id_last, seed_pos, { 0 }, is_dspark); for (int32_t i = 1; i < batch_len; ++i) { - common_batch_add(batch, mask_token_id, draft_pos_base + (i - 1), { 0 }, i <= n_keep); + common_batch_add(batch, mask_token_id, draft_pos_base + (i - 1), { 0 }, true); } if (llama_decode(ctx_dft, batch) != 0) { @@ -298,7 +316,8 @@ struct common_speculative_state_dflash : public common_speculative_state { for (int32_t i = 0; i < n_keep; ++i) { llama_token id = llama_get_dflash_draft_token_ith(ctx_dft, i); if (id == LLAMA_TOKEN_NULL) { - id = common_sampler_sample_speculative(nullptr, ctx_dft, i + 1, nullptr); + const int32_t logits_idx = is_dspark ? i : i + 1; + id = common_sampler_sample_speculative(nullptr, ctx_dft, logits_idx, nullptr); } result.push_back(id); } diff --git a/common/speculative.cpp b/common/speculative.cpp index 70ea82281..9783bd9fe 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -30,6 +30,7 @@ const std::vector common_speculative_types = { COMMON_SPECULATIVE_TYPE_NONE, COMMON_SPECULATIVE_TYPE_DRAFT, COMMON_SPECULATIVE_TYPE_DFLASH, + COMMON_SPECULATIVE_TYPE_DSPARK, COMMON_SPECULATIVE_TYPE_MTP, COMMON_SPECULATIVE_TYPE_EAGLE3, COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, @@ -44,6 +45,7 @@ const std::map common_speculative_typ {"none", COMMON_SPECULATIVE_TYPE_NONE}, {"draft", COMMON_SPECULATIVE_TYPE_DRAFT}, {"dflash", COMMON_SPECULATIVE_TYPE_DFLASH}, + {"dspark", COMMON_SPECULATIVE_TYPE_DSPARK}, {"mtp", COMMON_SPECULATIVE_TYPE_MTP}, {"eagle3", COMMON_SPECULATIVE_TYPE_EAGLE3}, {"ngram_simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE}, @@ -1210,6 +1212,7 @@ std::string common_speculative_type_to_str(enum common_speculative_type type) { case COMMON_SPECULATIVE_TYPE_NONE: return "none"; case COMMON_SPECULATIVE_TYPE_DRAFT: return "draft"; case COMMON_SPECULATIVE_TYPE_DFLASH: return "dflash"; + case COMMON_SPECULATIVE_TYPE_DSPARK: return "dspark"; case COMMON_SPECULATIVE_TYPE_MTP: return "mtp"; case COMMON_SPECULATIVE_TYPE_EAGLE3: return "eagle3"; case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram_simple"; @@ -1288,12 +1291,12 @@ common_speculative * common_speculative_init( } const bool has_dflash_stage = std::any_of(stages.begin(), stages.end(), [](const common_speculative_stage_params & stage) { - return stage.type == COMMON_SPECULATIVE_TYPE_DFLASH; + return common_speculative_type_is_dflash_family(stage.type); }); const bool needs_draft_ctx = std::any_of(stages.begin(), stages.end(), [¶ms](const common_speculative_stage_params & stage) { return stage.type == COMMON_SPECULATIVE_TYPE_DRAFT || - stage.type == COMMON_SPECULATIVE_TYPE_DFLASH || + common_speculative_type_is_dflash_family(stage.type) || (stage.type == COMMON_SPECULATIVE_TYPE_MTP && params.model_dft != nullptr); }); @@ -1314,7 +1317,7 @@ common_speculative * common_speculative_init( int32_t max_cross_ctx = 0; for (const auto & stage : stages) { - if (stage.type != COMMON_SPECULATIVE_TYPE_DFLASH) { + if (!common_speculative_type_is_dflash_family(stage.type)) { continue; } @@ -1401,14 +1404,16 @@ common_speculative * common_speculative_init( )); break; } - case COMMON_SPECULATIVE_TYPE_DFLASH: { + case COMMON_SPECULATIVE_TYPE_DFLASH: + case COMMON_SPECULATIVE_TYPE_DSPARK: { auto state = std::make_unique( config.type, ctx_tgt, ctx_dft, config.params.dflash_cross_ctx); if (!state->ready) { - LOG_ERR("%s: failed to initialize DFlash speculative state\n", __func__); + LOG_ERR("%s: failed to initialize %s speculative state\n", __func__, + common_speculative_type_to_str(config.type).c_str()); return nullptr; } impls.push_back(std::move(state)); @@ -1501,7 +1506,8 @@ common_speculative * common_speculative_init( } else if (params.autotune && !result->impls.empty()) { auto actual_type = result->impls[0]->type; if (actual_type != COMMON_SPECULATIVE_TYPE_NONE && - actual_type != COMMON_SPECULATIVE_TYPE_EAGLE3) { + actual_type != COMMON_SPECULATIVE_TYPE_EAGLE3 && + actual_type != COMMON_SPECULATIVE_TYPE_DSPARK) { result->tuner = std::make_unique(); result->tuner->init(actual_type, result->configs[0].params, llama_get_model(ctx_tgt)); LOG_DBG("Autotune initialized for %s, tuning %zu parameters\n", @@ -1666,6 +1672,16 @@ static bool common_speculative_has_type(const common_speculative * spec, common_ }); } +static bool common_speculative_has_dflash_family(const common_speculative * spec) { + if (spec == nullptr) { + return false; + } + + return std::any_of(spec->configs.begin(), spec->configs.end(), [](const common_speculative_config & config) { + return common_speculative_type_is_dflash_family(config.type); + }); +} + static int common_speculative_ctx_mtp_n_embd(llama_context * ctx) { return ctx ? (int) llama_mtp_state_n_embd(ctx) : 0; } @@ -1804,7 +1820,7 @@ static bool common_speculative_collect_target_batch_features( const llama_batch & batch, common_speculative_feature_view & features) { features = {}; - if (common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_DFLASH)) { + if (common_speculative_has_dflash_family(spec)) { return llama_spec_get_dflash_feature_view(ctx, batch, features); } @@ -1826,7 +1842,7 @@ static bool common_speculative_collect_target_seq_batch_features( llama_seq_id seq_id, common_speculative_feature_view & features) { features = {}; - if (common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_DFLASH)) { + if (common_speculative_has_dflash_family(spec)) { return llama_spec_get_dflash_feature_view_for_seq(ctx, batch, seq_id, features); } @@ -1912,8 +1928,13 @@ int common_speculative_get_configured_n_max(const common_speculative * spec) { } static bool common_speculative_has_target_features(const common_speculative * spec) { - return common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_MTP) || - common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_DFLASH); + if (spec == nullptr) { + return false; + } + + return std::any_of(spec->configs.begin(), spec->configs.end(), [](const common_speculative_config & config) { + return common_speculative_type_uses_target_features(config.type); + }); } bool common_speculative_load_draft_model( @@ -1958,7 +1979,7 @@ bool common_speculative_load_draft_model( if (params_dft.n_ctx == 0) { params_dft.n_ctx = params.n_ctx; } - if (params.has_stage_type(COMMON_SPECULATIVE_TYPE_DFLASH) && params_dft.n_gpu_layers < 0) { + if (params.has_dflash_family_stage() && params_dft.n_gpu_layers < 0) { params_dft.n_gpu_layers = params_base.n_gpu_layers; } params_dft.n_ctx = params_dft.n_ctx == 0 ? params_base.n_ctx / params_base.n_parallel : params_dft.n_ctx; @@ -2280,7 +2301,7 @@ bool common_speculative_copy_output_hidden_rows( const std::vector & output_indices, std::vector & hidden_rows) { hidden_rows.clear(); - if (common_speculative_has_type(spec, COMMON_SPECULATIVE_TYPE_DFLASH)) { + if (common_speculative_has_dflash_family(spec)) { return llama_spec_copy_dflash_rows_from_output_indices(ctx, output_indices, hidden_rows); } @@ -2706,7 +2727,7 @@ static common_speculative_state_dflash * common_speculative_get_dflash_state(com } for (auto & impl : spec->impls) { - if (impl->type != COMMON_SPECULATIVE_TYPE_DFLASH) { + if (!common_speculative_type_is_dflash_family(impl->type)) { continue; } diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 688193a36..55f879e92 100644 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -2623,6 +2623,66 @@ class DFlashDraftModel(Qwen3Model): return tensors +@Model.register("Qwen3DSparkModel") +class DSparkModel(DFlashDraftModel): + """Qwen3 DSpark sidecar: DFlash backbone plus a Markov head.""" + + model_arch = gguf.MODEL_ARCH.DFLASH_DRAFT + + _saw_markov_w1 = False + _saw_markov_w2 = False + _markov_shape: tuple[int, ...] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + top_level_name = name[6:] if name.startswith("model.") else name + + # DSpark shares vocabulary tensors with the exact target model supplied + # through --target-model-dir, matching the DFlash IO contract. + if top_level_name in ("embed_tokens.weight", "lm_head.weight"): + return [] + + if top_level_name in ("markov_head.markov_w1.weight", "markov_head.markov_w2.weight"): + if data_torch.ndim != 2: + raise ValueError(f"Qwen3DSparkModel: {top_level_name} must be rank-2, got {tuple(data_torch.shape)}") + shape = tuple(int(dim) for dim in data_torch.shape) + if self._markov_shape is None: + self._markov_shape = shape + elif shape != self._markov_shape: + raise ValueError( + "Qwen3DSparkModel: markov_w1 and markov_w2 shapes must match " + f"(got {self._markov_shape} and {shape})" + ) + tensor_id = ( + gguf.MODEL_TENSOR.DSPARK_MARKOV_W1 + if top_level_name.endswith("markov_w1.weight") + else gguf.MODEL_TENSOR.DSPARK_MARKOV_W2 + ) + if tensor_id == gguf.MODEL_TENSOR.DSPARK_MARKOV_W1: + self._saw_markov_w1 = True + else: + self._saw_markov_w2 = True + return [(f"{gguf.TENSOR_NAMES[tensor_id]}.weight", data_torch)] + + if top_level_name in ("confidence_head.proj.weight", "confidence_head.proj.bias"): + suffix = "bias" if top_level_name.endswith(".bias") else "weight" + return [(f"{gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.DSPARK_CONF_PROJ]}.{suffix}", data_torch)] + + return super().modify_tensors(data_torch, name, bid) + + def prepare_tensors(self): + super().prepare_tensors() + if not self._saw_markov_w1 or not self._saw_markov_w2: + raise ValueError( + "Qwen3DSparkModel conversion requires both " + "markov_head.markov_w1.weight and markov_head.markov_w2.weight" + ) + if self._markov_shape is None or self._markov_shape[0] != int(self.hparams.get("vocab_size", 0)): + raise ValueError( + "Qwen3DSparkModel: Markov tensors must have first dimension equal to " + f"vocab_size={self.hparams.get('vocab_size')} (shape={self._markov_shape})" + ) + + @Model.register("DFlashLagunaForCausalLM") class DFlashLagunaModel(DFlashDraftModel): model_arch = gguf.MODEL_ARCH.DFLASH_DRAFT diff --git a/examples/main/main.cpp b/examples/main/main.cpp index eca16500a..e4e700718 100644 --- a/examples/main/main.cpp +++ b/examples/main/main.cpp @@ -850,8 +850,7 @@ int main(int argc, char ** argv) { const bool need_prompt_target_features = embd_is_prompt && spec != nullptr && - (params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) || - params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_DFLASH)); + params.speculative.uses_target_features(); llama_batch batch = {}; if (need_prompt_target_features) { diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index 1c23f1708..620b67252 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -73,8 +73,7 @@ static bool server_response_needs_chat_parse(oaicompat_type oaicompat) { } static bool server_speculative_uses_target_features(const common_params_speculative & spec) { - return spec.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) || - spec.has_stage_type(COMMON_SPECULATIVE_TYPE_DFLASH); + return spec.uses_target_features(); } static bool server_speculative_requires_single_slot(const common_params_speculative & spec) { diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 630605cc5..e26e8f046 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -403,6 +403,9 @@ class MODEL_TENSOR(IntEnum): DFLASH_FC = auto() DFLASH_HIDDEN_NORM = auto() DFLASH_AUX_HIDDEN_NORM = auto() + DSPARK_MARKOV_W1 = auto() # DSpark Markov lookup matrix + DSPARK_MARKOV_W2 = auto() # DSpark Markov projection matrix + DSPARK_CONF_PROJ = auto() # DSpark confidence projection ATTN_KV = auto() ATTN_KV_NORM = auto() ATTN_OUT_A = auto() @@ -632,6 +635,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.DFLASH_FC: "dflash_fc", MODEL_TENSOR.DFLASH_HIDDEN_NORM: "dflash_hidden_norm", MODEL_TENSOR.DFLASH_AUX_HIDDEN_NORM: "dflash_aux_hidden_norm.{bid}", + MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", + MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", + MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", # openPangu-2.0 MODEL_TENSOR.INDEXER_K_NORM: "blk.{bid}.attn_indexer_k_norm", MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.attn_indexer_weights_proj", @@ -1560,6 +1566,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.DFLASH_FC, MODEL_TENSOR.DFLASH_HIDDEN_NORM, MODEL_TENSOR.DFLASH_AUX_HIDDEN_NORM, + # optional DSpark heads + MODEL_TENSOR.DSPARK_MARKOV_W1, + MODEL_TENSOR.DSPARK_MARKOV_W2, + MODEL_TENSOR.DSPARK_CONF_PROJ, ], MODEL_ARCH.BITNET: [ MODEL_TENSOR.ATTN_Q, diff --git a/src/graphs/build_deepseek4.cpp b/src/graphs/build_deepseek4.cpp index 5c62dc3b7..11346f7e2 100644 --- a/src/graphs/build_deepseek4.cpp +++ b/src/graphs/build_deepseek4.cpp @@ -8,6 +8,17 @@ #include #include +static ggml_tensor * dsv4_hc_mean_for_capture(ggml_context * ctx, ggml_tensor * x) { + GGML_ASSERT(x != nullptr && x->ne[1] > 0); + + ggml_tensor * mean = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0); + for (int64_t hc_row = 1; hc_row < x->ne[1]; ++hc_row) { + mean = ggml_add(ctx, mean, + ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], hc_row * x->nb[1])); + } + return ggml_scale(ctx, mean, 1.0f / (float) x->ne[1]); +} + static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) { if (ext_factor == 0.0f) { return 1.0f; @@ -1407,7 +1418,13 @@ ggml_cgraph * llm_build_context::build_deepseek4() { inpL = build_mhc_post(cur, post, residual, comb, n_embd, hc, true); inpL = lctx.cvec.apply_to(ctx0, inpL, il); - cb(inpL, "l_out", il); + if (lctx.dflash.capture) { + ggml_tensor * capture = dsv4_hc_mean_for_capture(ctx0, inpL); + cb(capture, "l_out", il); + ggml_build_forward_expand(gf, capture); + } else { + cb(inpL, "l_out", il); + } } if (is_mtp) { @@ -1472,3 +1489,179 @@ ggml_cgraph * llm_build_context::build_deepseek4() { return gf; } + +ggml_cgraph * llm_build_context::build_dflash_dsv4() { + const int64_t n_embd_head = hparams.n_embd_head_k(0); + const int64_t n_embd_head_rope = hparams.n_rot; + const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; + const int64_t ctx_len = lctx.dflash.visible_cross_ctx > 0 + ? (int64_t) lctx.dflash.visible_cross_ctx + : std::max(1, (int64_t) cparams.n_ctx - (int64_t) hparams.dflash_block_size); + const int64_t n_kv_total = GGML_PAD(ctx_len + n_tokens, flash_attn ? 256 : 32); + const ggml_type mask_type = flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; + + if (!lctx.ensure_dflash_kv_cache_tensors((int32_t) ctx_len)) { + LLAMA_LOG_ERROR("%s: failed to initialize DSV4 DFlash K/V cache\n", __func__); + return nullptr; + } + + GGML_ASSERT(n_embd_head_nope > 0); + GGML_ASSERT(hparams.n_head_kv() == 1); + GGML_ASSERT(model.dflash_fc != nullptr); + GGML_ASSERT(model.dflash_hidden_norm != nullptr); + GGML_ASSERT(model.hc_head_fn != nullptr && model.hc_head_base != nullptr && model.hc_head_scale != nullptr); + + ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes((int) std::max(n_tokens, ctx_len)) + 48 * n_layer, false); + + lctx.dflash.inputs.kq_mask = nullptr; + lctx.dflash.kv.kq_mask_tensor = nullptr; + lctx.dflash.inputs.kq_mask_swa = ggml_new_tensor_2d(ctx0, mask_type, n_kv_total, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)); + lctx.dflash.kv.kq_mask_swa_tensor = lctx.dflash.inputs.kq_mask_swa; + ggml_set_input(lctx.dflash.inputs.kq_mask_swa); + cb(lctx.dflash.inputs.kq_mask_swa, "dsv4_dflash_kq_mask_swa", -1); + + lctx.dflash.kv.draft_tail_rows_tensor = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(lctx.dflash.kv.draft_tail_rows_tensor); + cb(lctx.dflash.kv.draft_tail_rows_tensor, "dflash_draft_tail_rows", -1); + + ggml_tensor * tok_embd = model.tok_embd; + GGML_ASSERT(tok_embd != nullptr); + ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, tok_embd, cb); + inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens); + inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hparams.dsv4_hc_mult, n_tokens, 1); + cb(inpL, "dsv4_dflash_hc_init", -1); + + auto build_attention = [&](int il, ggml_tensor * input, ggml_tensor * inp_pos) { + const auto & layer = model.layers[il]; + ggml_tensor * qr = llm_build_lora_mm(lctx, ctx0, layer.wq_a, input); + qr = llm_build_norm(ctx0, qr, hparams, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, cb, il); + ggml_tensor * q = llm_build_lora_mm(lctx, ctx0, layer.wq_b, qr); + q = ggml_reshape_2d(ctx0, q, n_embd_head, n_head * n_tokens); + q = ggml_rms_norm(ctx0, q, hparams.f_norm_rms_eps); + q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, n_tokens); + q = ggml_rope_ext_inplace(ctx0, q, inp_pos, nullptr, n_embd_head_rope, rope_type, 0, + freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + q->op_params[15] = 1; + cb(q, "dsv4_dflash_q", il); + + ggml_tensor * kv = llm_build_lora_mm(lctx, ctx0, layer.wkv_latent, input); + kv = llm_build_norm(ctx0, kv, hparams, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, cb, il); + kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens); + + kv = ggml_rope_ext_inplace(ctx0, kv, inp_pos, nullptr, n_embd_head_rope, rope_type, 0, + freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + kv->op_params[15] = 1; + kv = ggml_permute(ctx0, kv, 0, 2, 1, 3); + cb(kv, "dsv4_dflash_kv", il); + + ggml_tensor * Kcur = ggml_set_rows(ctx0, lctx.dflash.kv.k_ctx_cache[il], kv, + lctx.dflash.kv.draft_tail_rows_tensor); + ggml_tensor * Vcur = ggml_set_rows(ctx0, lctx.dflash.kv.v_ctx_cache[il], kv, + lctx.dflash.kv.draft_tail_rows_tensor); + cb(Kcur, "dsv4_dflash_k_set_tail", il); + cb(Vcur, "dsv4_dflash_v_set_tail", il); + ggml_build_forward_expand(gf, Kcur); + ggml_build_forward_expand(gf, Vcur); + Kcur = ggml_view_3d(ctx0, lctx.dflash.kv.k_ctx_cache[il], + lctx.dflash.kv.k_ctx_cache[il]->ne[0], n_kv_total, + lctx.dflash.kv.k_ctx_cache[il]->ne[2], + lctx.dflash.kv.k_ctx_cache[il]->nb[1], + lctx.dflash.kv.k_ctx_cache[il]->nb[2], 0); + Vcur = ggml_view_3d(ctx0, lctx.dflash.kv.v_ctx_cache[il], + lctx.dflash.kv.v_ctx_cache[il]->ne[0], n_kv_total, + lctx.dflash.kv.v_ctx_cache[il]->ne[2], + lctx.dflash.kv.v_ctx_cache[il]->nb[1], + lctx.dflash.kv.v_ctx_cache[il]->nb[2], 0); + if (Kcur->type == GGML_TYPE_F32) { Kcur = ggml_cast(ctx0, Kcur, GGML_TYPE_F16); } + if (Vcur->type == GGML_TYPE_F32) { Vcur = ggml_cast(ctx0, Vcur, GGML_TYPE_F16); } + + ggml_tensor * q_attn = ggml_permute(ctx0, q, 0, 2, 1, 3); + ggml_tensor * mask = lctx.dflash.inputs.kq_mask_swa; + ggml_tensor * attn = ggml_flash_attn_ext(ctx0, q_attn, Kcur, Vcur, mask, + 1.0f / std::sqrt((float) n_embd_head), hparams.f_max_alibi_bias, + hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); + if (layer.attn_sinks) { + ggml_flash_attn_ext_add_sinks(attn, layer.attn_sinks); + } + cb(attn, "dsv4_dflash_attn", il); + ggml_build_forward_expand(gf, attn); + + attn = ggml_reshape_3d(ctx0, attn, n_embd_head, n_head, n_tokens); + attn = ggml_rope_ext_inplace(ctx0, attn, inp_pos, nullptr, n_embd_head_rope, rope_type, 0, + freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + attn->op = GGML_OP_ROPE_BACK; + attn->op_params[15] = 1; + + const int64_t o_group_dim = layer.wo_a->ne[0]; + const int64_t n_groups = (n_head * n_embd_head) / o_group_dim; + const int64_t o_lora_rank = layer.wo_b->ne[0] / n_groups; + attn = ggml_reshape_3d(ctx0, attn, o_group_dim, n_groups, n_tokens); + attn = ggml_permute(ctx0, attn, 0, 2, 1, 3); + ggml_tensor * oa = ggml_mul_mat(ctx0, + ggml_reshape_3d(ctx0, layer.wo_a, layer.wo_a->ne[0], o_lora_rank, n_groups), attn); + oa = ggml_permute(ctx0, oa, 0, 2, 1, 3); + oa = ggml_cont_2d(ctx0, oa, o_lora_rank * n_groups, n_tokens); + return llm_build_lora_mm(lctx, ctx0, layer.wo_b, oa); + }; + + ggml_tensor * inp_pos = build_inp_pos(); + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + ggml_tensor * cur = build_hc_pre(ctx0, *this, hparams, n_embd, hparams.f_norm_rms_eps, inpL, + layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, cb, il); + cur = llm_build_norm(ctx0, cur, hparams, layer.attn_norm, nullptr, LLM_NORM_RMS, cb, il); + cur = build_attention(il, cur, inp_pos); + inpL = build_mhc_post(cur, post, residual, comb, n_embd, hparams.dsv4_hc_mult, true); + + residual = inpL; + cur = build_hc_pre(ctx0, *this, hparams, n_embd, hparams.f_norm_rms_eps, inpL, + layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, cb, il); + cur = llm_build_norm(ctx0, cur, hparams, layer.ffn_norm, nullptr, LLM_NORM_RMS, cb, il); + ggml_tensor * moe = llm_build_moe_ffn(ctx0, lctx, cur, + layer.ffn_gate_inp, nullptr, + layer.ffn_up_exps, nullptr, + layer.ffn_gate_exps, nullptr, + layer.ffn_down_exps, nullptr, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, true, hparams.expert_weights_scale, + (llm_expert_gating_func_type) hparams.expert_gating_func, + cb, il, gf, false, layer.ffn_up_gate_exps, nullptr, nullptr, nullptr, nullptr); + ggml_tensor * shexp = llm_build_ffn(ctx0, lctx, nullptr, cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf, false, false); + cur = ggml_add(ctx0, moe, shexp); + inpL = build_mhc_post(cur, post, residual, comb, n_embd, hparams.dsv4_hc_mult, true); + } + + ggml_tensor * out = build_hc_head(ctx0, *this, hparams, n_embd, hparams.f_norm_rms_eps, + inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base); + out = llm_build_norm(ctx0, out, hparams, model.output_norm, nullptr, LLM_NORM_RMS, cb, -1); + out = build_output(lctx, ctx0, out, model.output, nullptr, cb); + if (lctx.dflash.dspark) { + cb(out, "dflash_base_result_output", -1); + } else { + cb(out, "result_output", -1); + } + ggml_build_forward_expand(gf, out); + + lctx.dflash.draft_tokens_tensor = nullptr; + ggml_tensor * draft_tokens = nullptr; + ggml_tensor * result = out; + if (lctx.dflash.dspark) { + result = build_dspark_logits(*this, out, lctx.inp_tokens, &draft_tokens); + cb(result, "result_output", -1); + } else { + draft_tokens = ggml_argmax(ctx0, result); + } + ggml_set_name(draft_tokens, "draft_argmax"); + ggml_build_forward_expand(gf, result); + ggml_build_forward_expand(gf, draft_tokens); + lctx.dflash.draft_tokens_tensor = draft_tokens; + return gf; +} diff --git a/src/graphs/build_dflash.cpp b/src/graphs/build_dflash.cpp index e7dfd5bba..821b5c039 100644 --- a/src/graphs/build_dflash.cpp +++ b/src/graphs/build_dflash.cpp @@ -4,6 +4,53 @@ #include +ggml_tensor * llm_build_context::build_dspark_logits( + llm_build_context & llm, + ggml_tensor * base_logits, + ggml_tensor * input_tokens, + ggml_tensor ** draft_tokens_out) { + ggml_context * ctx0 = llm.ctx0; + const llama_model & model = llm.model; + + GGML_ASSERT(model.dspark_markov_w1 != nullptr); + GGML_ASSERT(model.dspark_markov_w2 != nullptr); + GGML_ASSERT(input_tokens != nullptr); + GGML_ASSERT(base_logits->ne[1] == input_tokens->ne[0]); + + const int64_t n_vocab = base_logits->ne[0]; + const int64_t n_tokens = base_logits->ne[1]; + GGML_ASSERT(n_tokens > 0); + + ggml_tensor * previous = ggml_view_1d(ctx0, input_tokens, 1, 0); + ggml_tensor * chained = nullptr; + ggml_tensor * draft_tokens = nullptr; + + for (int64_t i = 0; i < n_tokens; ++i) { + ggml_tensor * markov_w1 = ggml_get_rows(ctx0, model.dspark_markov_w1, previous); + ggml_tensor * markov_bias = ggml_mul_mat(ctx0, model.dspark_markov_w2, markov_w1); + ggml_tensor * base_row = ggml_view_2d( + ctx0, + base_logits, + n_vocab, + 1, + base_logits->nb[1], + (size_t) i * base_logits->nb[1]); + ggml_tensor * biased_row = ggml_add(ctx0, base_row, markov_bias); + ggml_tensor * token = ggml_argmax(ctx0, biased_row); + + chained = chained == nullptr ? biased_row : ggml_concat(ctx0, chained, biased_row, 1); + draft_tokens = draft_tokens == nullptr ? token : ggml_concat(ctx0, draft_tokens, token, 0); + if (i + 1 < n_tokens) { + previous = token; + } + } + + if (draft_tokens_out != nullptr) { + *draft_tokens_out = draft_tokens; + } + return chained; +} + ggml_cgraph * llm_build_context::build_dflash_kv_cache() { const int64_t n_embd_head_k = hparams.n_embd_head_k(0); const int64_t n_embd_head_v = hparams.n_embd_head_v(0); @@ -30,6 +77,14 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { ggml_set_input(lctx.dflash.kv.cache_input_pos_ctx); cb(lctx.dflash.kv.cache_input_pos_ctx, "dflash_kv_input_pos_ctx", -1); + if (hparams.dflash_dsv4) { + lctx.dflash.kv.cache_input_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, update_rows); + ggml_set_input(lctx.dflash.kv.cache_input_rows); + cb(lctx.dflash.kv.cache_input_rows, "dflash_kv_input_rows", -1); + } else { + lctx.dflash.kv.cache_input_rows = nullptr; + } + ggml_tensor * target_features = lctx.dflash.kv.cache_input_target_features; if (hparams.dflash_laguna) { GGML_ASSERT(model.dflash_aux_hidden_norms.size() == hparams.dflash_n_target_layers); @@ -66,6 +121,44 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { fused_target = llm_build_norm(ctx0, fused_target, hparams, model.dflash_hidden_norm, nullptr, LLM_NORM_RMS, cb, -1); cb(fused_target, "dflash_kv_fused_target", -1); + if (hparams.dflash_dsv4) { + const int64_t n_embd_head = hparams.n_embd_head_k(0); + const int64_t n_embd_head_rope = hparams.n_rot; + const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; + + GGML_ASSERT(n_embd_head_nope > 0); + GGML_ASSERT(hparams.n_head_kv() == 1); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + GGML_ASSERT(layer.wkv_latent != nullptr); + GGML_ASSERT(layer.attn_kv_norm != nullptr); + + ggml_tensor * kv = llm_build_lora_mm(lctx, ctx0, layer.wkv_latent, fused_target); + kv = llm_build_norm(ctx0, kv, hparams, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, cb, il); + kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, update_rows); + + kv = ggml_rope_ext_inplace(ctx0, kv, lctx.dflash.kv.cache_input_pos_ctx, nullptr, + n_embd_head_rope, rope_type, 0, freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + kv->op_params[15] = 1; + cb(kv, "dflash_dsv4_kv_injected", il); + kv = ggml_permute(ctx0, kv, 0, 2, 1, 3); + + GGML_ASSERT(lctx.dflash.kv.k_ctx_cache[il] != nullptr); + GGML_ASSERT(lctx.dflash.kv.v_ctx_cache[il] != nullptr); + ggml_tensor * Kcur = ggml_set_rows(ctx0, lctx.dflash.kv.k_ctx_cache[il], kv, + lctx.dflash.kv.cache_input_rows); + ggml_tensor * Vcur = ggml_set_rows(ctx0, lctx.dflash.kv.v_ctx_cache[il], kv, + lctx.dflash.kv.cache_input_rows); + cb(Kcur, "dflash_dsv4_k_set_rows", il); + cb(Vcur, "dflash_dsv4_v_set_rows", il); + ggml_build_forward_expand(gf, Kcur); + ggml_build_forward_expand(gf, Vcur); + } + + return gf; + } + for (int il = 0; il < n_layer; ++il) { GGML_ASSERT(il < (int32_t) lctx.dflash.kv.k_ctx_cache.size()); GGML_ASSERT(il < (int32_t) lctx.dflash.kv.v_ctx_cache.size()); @@ -201,6 +294,10 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { } ggml_cgraph * llm_build_context::build_dflash() { + if (hparams.dflash_dsv4) { + return build_dflash_dsv4(); + } + const int64_t n_embd_head_k = hparams.n_embd_head_k(0); const int64_t n_embd_head_v = hparams.n_embd_head_v(0); const int64_t n_target_features = hparams.dflash_n_target_features; @@ -432,12 +529,23 @@ ggml_cgraph * llm_build_context::build_dflash() { GGML_ASSERT(model.output_mtp != nullptr); ggml_tensor * result = build_output(lctx, ctx0, inpL, model.output_mtp, model.output_norm, cb); - cb(result, "result_output", -1); + if (lctx.dflash.dspark) { + cb(result, "dflash_base_result_output", -1); + } else { + cb(result, "result_output", -1); + } ggml_build_forward_expand(gf, result); lctx.dflash.draft_tokens_tensor = nullptr; - ggml_tensor * draft_tokens = ggml_argmax(ctx0, result); + ggml_tensor * draft_tokens = nullptr; + if (lctx.dflash.dspark) { + result = build_dspark_logits(*this, result, lctx.inp_tokens, &draft_tokens); + cb(result, "result_output", -1); + } else { + draft_tokens = ggml_argmax(ctx0, result); + } ggml_set_name(draft_tokens, "draft_argmax"); + ggml_build_forward_expand(gf, result); ggml_build_forward_expand(gf, draft_tokens); lctx.dflash.draft_tokens_tensor = draft_tokens; diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 9c05d3b8b..dfb83aaf2 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -85,6 +85,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MISTRAL4, "mistral4" }, { LLM_ARCH_GEMMA4, "gemma4" }, { LLM_ARCH_GEMMA4_MTP, "gemma4_mtp" }, + { LLM_ARCH_DFLASH, "dflash" }, { LLM_ARCH_DFLASH_DRAFT, "dflash-draft" }, { LLM_ARCH_GEMMA4_ASSISTANT,"gemma4-assistant" }, { LLM_ARCH_OPENPANGU, "openpangu" }, @@ -167,6 +168,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_DFLASH_BLOCK_SIZE, "%s.dflash.block_size" }, { LLM_KV_DFLASH_MASK_TOKEN_ID, "%s.dflash.mask_token_id" }, { LLM_KV_DFLASH_TARGET_LAYER_IDS, "%s.dflash.target_layer_ids" }, + { LLM_KV_DFLASH_TARGET_LAYERS, "%s.target_layers" }, { LLM_KV_DFLASH_N_TARGET_FEATURES, "%s.dflash.n_target_features" }, { LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, "%s.dflash.backbone_rotary_base" }, { LLM_KV_DFLASH_LAGUNA, "%s.dflash.laguna" }, @@ -318,3 +320,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { return false; } } + +bool llm_arch_is_dflash_family(const llm_arch & arch) { + return arch == LLM_ARCH_DFLASH || arch == LLM_ARCH_DFLASH_DRAFT; +} diff --git a/src/llama-arch.h b/src/llama-arch.h index 2f4224a34..b8687d518 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -83,6 +83,7 @@ enum llm_arch { LLM_ARCH_MISTRAL4, LLM_ARCH_GEMMA4, LLM_ARCH_GEMMA4_MTP, + LLM_ARCH_DFLASH, LLM_ARCH_DFLASH_DRAFT, LLM_ARCH_GEMMA4_ASSISTANT, LLM_ARCH_OPENPANGU, @@ -150,6 +151,7 @@ enum llm_kv { LLM_KV_DFLASH_BLOCK_SIZE, LLM_KV_DFLASH_MASK_TOKEN_ID, LLM_KV_DFLASH_TARGET_LAYER_IDS, + LLM_KV_DFLASH_TARGET_LAYERS, LLM_KV_DFLASH_N_TARGET_FEATURES, LLM_KV_DFLASH_BACKBONE_ROTARY_BASE, LLM_KV_DFLASH_LAGUNA, @@ -421,6 +423,9 @@ enum llm_tensor { LLM_TENSOR_DFLASH_FC, LLM_TENSOR_DFLASH_HIDDEN_NORM, LLM_TENSOR_DFLASH_AUX_HIDDEN_NORM, + LLM_TENSOR_DSPARK_MARKOV_W1, + LLM_TENSOR_DSPARK_MARKOV_W2, + LLM_TENSOR_DSPARK_CONF_PROJ, // openPangu-2.0 LLM_TENSOR_ATTN_QA_CONV, // MoME causal conv on q-lora latent @@ -451,5 +456,6 @@ const char * llama_model_arch_name(llm_arch arch); bool llm_arch_is_recurrent(const llm_arch & arch); bool llm_arch_is_hybrid(const llm_arch & arch); +bool llm_arch_is_dflash_family(const llm_arch & arch); llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int il); diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 6eb935d10..a7f6067b2 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -2823,6 +2823,7 @@ ggml_cgraph * llm_build_context::llama_build_graph( { result = llm.build_gemma4_mtp(); } break; + case LLM_ARCH_DFLASH: case LLM_ARCH_DFLASH_DRAFT: { result = llm.build_dflash(); diff --git a/src/llama-build-context.h b/src/llama-build-context.h index dc851ade8..ba489eea6 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -272,6 +272,8 @@ struct llm_build_context { ggml_cgraph * build_dflash(); + ggml_cgraph * build_dflash_dsv4(); + ggml_cgraph * build_dflash_kv_cache(); ggml_cgraph * build_starcoder2(); @@ -505,6 +507,10 @@ struct llm_build_context { const llm_build_cb & cb, int il, ggml_cgraph * graph = nullptr, bool add_input = false, bool is_norm = false, ggml_tensor * add_extra = nullptr, ggml_tensor * post_norm = nullptr); + static ggml_tensor * build_dspark_logits(llm_build_context & llm, + ggml_tensor * base_logits, ggml_tensor * input_tokens, + ggml_tensor ** draft_tokens = nullptr); + static ggml_tensor * llm_build_moe_ffn(ggml_context * ctx, llama_context & lctx, ggml_tensor * cur, ggml_tensor * gate_inp, ggml_tensor * gate_inp_b, diff --git a/src/llama-context.h b/src/llama-context.h index 4581e5a30..5b296f61b 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -439,6 +439,7 @@ struct llama_context { int32_t cache_graph_write_pos = 0; struct ggml_tensor * cache_input_target_features = nullptr; struct ggml_tensor * cache_input_pos_ctx = nullptr; + struct ggml_tensor * cache_input_rows = nullptr; struct ggml_tensor * kq_mask_tensor = nullptr; struct ggml_tensor * kq_mask_swa_tensor = nullptr; struct ggml_tensor * draft_tail_rows_tensor = nullptr; @@ -472,9 +473,8 @@ struct llama_context { std::vector feature_view_buffer; input_state inputs; int32_t visible_cross_ctx = 0; + bool dspark = false; - // Argmax token IDs from the DFlash draft graph, computed via GPU argmax. - // Populated in llama_decode_internal after graph compute. std::vector draft_tokens; struct ggml_tensor * draft_tokens_tensor = nullptr; }; diff --git a/src/llama-dflash.cpp b/src/llama-dflash.cpp index 1d36dbef5..de45845fb 100644 --- a/src/llama-dflash.cpp +++ b/src/llama-dflash.cpp @@ -182,6 +182,7 @@ void llama_context::free_dflash_kv_cache_tensors() { dflash.kv.cache_graph_write_pos = 0; dflash.kv.cache_input_target_features = nullptr; dflash.kv.cache_input_pos_ctx = nullptr; + dflash.kv.cache_input_rows = nullptr; dflash.kv.kq_mask_tensor = nullptr; dflash.kv.kq_mask_swa_tensor = nullptr; dflash.kv.draft_tail_rows_tensor = nullptr; @@ -324,6 +325,21 @@ static bool validate_dflash_graph_contract(const llama_context & lctx) { return false; } + if (hparams.dflash_dsv4) { + if (model.layers[il].wq_a == nullptr || + model.layers[il].wq_b == nullptr || + model.layers[il].attn_q_a_norm == nullptr || + model.layers[il].wkv_latent == nullptr || + model.layers[il].attn_kv_norm == nullptr || + model.layers[il].wo_a == nullptr || + model.layers[il].wo_b == nullptr) { + LLAMA_LOG_ERROR("%s: DSV4 DFlash layer %d is missing a required Q/KV/output projection tensor\n", + __func__, il); + return false; + } + continue; + } + if (model.layers[il].attn_norm == nullptr || model.layers[il].attn_q_norm == nullptr || model.layers[il].attn_k_norm == nullptr) { @@ -364,6 +380,7 @@ bool llama_prepare_dflash_graph_inputs( const int32_t cross_ctx = lctx.dflash.visible_cross_ctx > 0 ? lctx.dflash.visible_cross_ctx : std::max(1, (int32_t) lctx.cparams.n_ctx - (int32_t) lctx.model.hparams.dflash_block_size); + const bool is_dsv4 = lctx.model.hparams.dflash_dsv4; ggml_tensor * kq_mask = lctx.dflash.kv.kq_mask_tensor; ggml_tensor * kq_mask_swa = lctx.dflash.kv.kq_mask_swa_tensor; @@ -528,12 +545,15 @@ bool llama_prepare_dflash_graph_inputs( ggml_cgraph * gf_kv = nullptr; const bool can_reuse_kv_graph = lctx.dflash.kv.cache_graph != nullptr && lctx.dflash.kv.cache_graph_rows == update_rows && - lctx.dflash.kv.cache_graph_write_pos == lctx.dflash.kv.cache_write_pos; + (lctx.model.hparams.dflash_dsv4 || + lctx.dflash.kv.cache_graph_write_pos == lctx.dflash.kv.cache_write_pos); if (can_reuse_kv_graph) { gf_kv = lctx.dflash.kv.cache_graph; } else { gf_kv = llm_build_context::llama_build_graph_dflash_kv_cache(lctx); - if (gf_kv == nullptr || lctx.dflash.kv.cache_input_target_features == nullptr || lctx.dflash.kv.cache_input_pos_ctx == nullptr) { + if (gf_kv == nullptr || lctx.dflash.kv.cache_input_target_features == nullptr || + lctx.dflash.kv.cache_input_pos_ctx == nullptr || + (lctx.model.hparams.dflash_dsv4 && lctx.dflash.kv.cache_input_rows == nullptr)) { LLAMA_LOG_ERROR("%s: failed to build DFlash K/V cache graph\n", __func__); return false; } @@ -559,6 +579,20 @@ bool llama_prepare_dflash_graph_inputs( } else { ggml_backend_tensor_set(lctx.dflash.kv.cache_input_pos_ctx, update_pos, 0, ggml_nbytes(lctx.dflash.kv.cache_input_pos_ctx)); } + if (lctx.model.hparams.dflash_dsv4) { + std::vector update_rows_idx((size_t) update_rows); + for (int32_t i = 0; i < update_rows; ++i) { + update_rows_idx[(size_t) i] = (cache_write_start + i) % cross_ctx; + } + ggml_backend_t kv_rows_backend = llama_backend_for_tensor(lctx, lctx.dflash.kv.cache_input_rows); + if (kv_rows_backend != nullptr) { + ggml_backend_tensor_set_async(kv_rows_backend, lctx.dflash.kv.cache_input_rows, + update_rows_idx.data(), 0, ggml_nbytes(lctx.dflash.kv.cache_input_rows)); + } else { + ggml_backend_tensor_set(lctx.dflash.kv.cache_input_rows, update_rows_idx.data(), 0, + ggml_nbytes(lctx.dflash.kv.cache_input_rows)); + } + } llama_graph_compute_sched(lctx, lctx.dflash.kv.cache_sched, gf_kv, lctx.cparams.n_threads); ggml_backend_sched_synchronize(lctx.dflash.kv.cache_sched); @@ -659,10 +693,8 @@ bool llama_prepare_dflash_graph_inputs( for (int32_t k = cross_ctx; k < cross_ctx + (int32_t) n_tokens; ++k) { const int32_t block_k = k - cross_ctx; - // intra-block draft tokens are contiguous from draft_pos_base, so the - // SWA distance is (j - block_k); apply the same window bound as the - // cross-context section above (causal AND within n_swa). - if (block_k <= (int32_t) j && ((int32_t) j - block_k) < swa_window) { + // DSV4 Dspark rows see the complete current block, standard DFlash is causal. + if ((is_dsv4 || block_k <= (int32_t) j) && ((int32_t) j - block_k) < swa_window) { row[k] = h_zero; } } @@ -689,7 +721,7 @@ bool llama_prepare_dflash_graph_inputs( // intra-block draft tokens are contiguous from draft_pos_base, so the // SWA distance is (j - block_k); apply the same window bound as the // cross-context section above (causal AND within n_swa). - if (block_k <= (int32_t) j && ((int32_t) j - block_k) < swa_window) { + if ((is_dsv4 || block_k <= (int32_t) j) && ((int32_t) j - block_k) < swa_window) { row[k] = 0.0f; } } diff --git a/src/llama-dsv4.cpp b/src/llama-dsv4.cpp index 59db88ea2..8e87e08ba 100644 --- a/src/llama-dsv4.cpp +++ b/src/llama-dsv4.cpp @@ -1226,10 +1226,19 @@ static bool dsv4_per_step_copy_base(llama_context & ctx, bool restore) { for (size_t i = 0; i < ckpt.dsv4_per_step_state.size(); ++i) { ggml_tensor * state = ckpt.dsv4_per_step_state[i]; ggml_tensor * shadow = ckpt.dsv4_per_step_state_shadow[i]; - ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(ctx.sched, state); - if (state == nullptr || shadow == nullptr || backend == nullptr) { + ggml_backend_t backend = state != nullptr + ? ggml_backend_sched_get_tensor_backend(ctx.sched, state) + : nullptr; + if (state == nullptr || shadow == nullptr) { return false; } + if (backend == nullptr) { + if (state->buffer == nullptr || shadow->buffer == nullptr) { + return false; + } + ggml_backend_tensor_copy(restore ? shadow : state, restore ? state : shadow); + continue; + } if (restore) { ggml_backend_tensor_copy_async(backend, backend, shadow, state); } else { @@ -1502,7 +1511,7 @@ static enum llama_spec_ckpt_restore_result dsv4_per_step_restore_rows( return LLAMA_SPEC_CKPT_RESTORE_FAILED; } ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(ctx.sched, state); - if (backend == nullptr) { + if (backend == nullptr && (state->buffer == nullptr || delta->buffer == nullptr)) { return LLAMA_SPEC_CKPT_RESTORE_FAILED; } @@ -1529,9 +1538,13 @@ static enum llama_spec_ckpt_restore_result dsv4_per_step_restore_rows( dst_view.view_src = nullptr; src_view.view_offs = 0; dst_view.view_offs = 0; - ggml_backend_tensor_copy_async(backend, backend, &src_view, &dst_view); + if (backend != nullptr) { + ggml_backend_tensor_copy_async(backend, backend, &src_view, &dst_view); + } else { + ggml_backend_tensor_copy(&src_view, &dst_view); + } } - if (std::find(backends.begin(), backends.end(), backend) == backends.end()) { + if (backend != nullptr && std::find(backends.begin(), backends.end(), backend) == backends.end()) { backends.push_back(backend); } } diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index f528710ed..81eac9e9b 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -1673,9 +1673,11 @@ void llm_load_hparams( default: model.type = e_model::MODEL_UNKNOWN; } } break; + case LLM_ARCH_DFLASH: case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM_DSA: { + const bool is_dsv4 = model.arch == LLM_ARCH_DFLASH || model.arch == LLM_ARCH_DEEPSEEK4; ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false); if (model.arch == LLM_ARCH_DEEPSEEK4 && hparams.n_layer == 43 && hparams.nextn_predict_layers > 0) { LLAMA_LOG_WARN("===============================================================================================\n"); @@ -1710,7 +1712,7 @@ void llm_load_hparams( ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); - if (model.arch == LLM_ARCH_DEEPSEEK4) { + if (is_dsv4) { ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_limits, hparams.n_layer); if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_limits_shared, hparams.n_layer, 0)) { hparams.swiglu_limits_shared = hparams.swiglu_limits; @@ -1722,7 +1724,7 @@ void llm_load_hparams( // this does not select the DeepSeek V3 MLA path. ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv, false); - if (model.arch == LLM_ARCH_DEEPSEEK4 && hparams.n_lora_kv == 0) { + if (is_dsv4 && hparams.n_lora_kv == 0) { const uint32_t probe_layer = dsv4_probe_layer; if (auto * kv_norm = ml.get_tensor_meta(format("blk.%u.attn_kv_a_norm.weight", probe_layer).c_str())) { hparams.n_lora_kv = (uint32_t) kv_norm->ne[0]; @@ -1763,7 +1765,7 @@ void llm_load_hparams( ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); - if (model.arch == LLM_ARCH_DEEPSEEK4) { + if (is_dsv4) { ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); hparams.rope_freq_scale_train_swa = 1.0f; @@ -1858,12 +1860,12 @@ void llm_load_hparams( hparams.expert_gating_func = LLM_EXPERT_GATING_FUNC_SIGMOID; } - if (model.arch == LLM_ARCH_DEEPSEEK4 && + if (is_dsv4 && hparams.expert_gating_func != LLM_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { throw std::runtime_error("DeepSeek-V4 loader currently expects sqrtsoftplus MoE scoring"); } - if (model.arch == LLM_ARCH_DEEPSEEK4) { + if (is_dsv4) { hparams.n_layer_kv_from_start = hparams.n_layer; } else { ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false); @@ -1877,7 +1879,7 @@ void llm_load_hparams( case 79: model.type = MODEL_744B_A40B; break; default: model.type = MODEL_UNKNOWN; } - if (model.arch != LLM_ARCH_DEEPSEEK4 && hparams.n_head_kv() == 1) { + if (!is_dsv4 && hparams.n_head_kv() == 1) { int n_nead_kv = hparams.n_gqa(); if (n_nead_kv%4 != 0 || hparams.n_embd_head_k_full != 576 || hparams.n_embd_head_v_full != 512 || hparams.n_rot != 64) { @@ -1894,6 +1896,37 @@ void llm_load_hparams( ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_full); ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_full); } + + if (model.arch == LLM_ARCH_DFLASH) { + hparams.dflash_dsv4 = hparams.dsv4_hc_mult > 0; + if (!hparams.dflash_dsv4) { + throw std::runtime_error("dflash: hyper_connection.count is required for the official DSV4 schema"); + } + + ml.get_key("dflash.block_size", hparams.dflash_block_size, true); + ml.get_key(LLM_KV_TOKENIZER_MASK_ID, hparams.dflash_mask_token_id, true); + load_dflash_target_layer_ids( + ml, + LLM_KV(model.arch)(LLM_KV_DFLASH_TARGET_LAYERS), + hparams, + true); + + for (uint32_t i = 0; i < hparams.dflash_n_target_layers; ++i) { + if (hparams.dflash_target_layer_ids[i] == 0) { + throw std::runtime_error("dflash: target_layers must use one-based IDs"); + } + --hparams.dflash_target_layer_ids[i]; + } + + hparams.dflash_n_target_features = hparams.n_embd * hparams.dflash_n_target_layers; + hparams.dflash_laguna = false; + for (uint32_t i = 0; i < hparams.n_layer; ++i) { + if (hparams.dsv4_compress_ratios[i] != 0) { + throw std::runtime_error("dflash: DSV4 draft requires uncompressed stages"); + } + } + validate_dflash_hparams(hparams, model.arch); + } } break; default: (void)0; } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 905c0598a..cc1d43db7 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -170,6 +170,7 @@ struct llama_hparams { uint32_t dflash_target_layer_ids[8] = {}; float dflash_backbone_rotary_base = 0.0f; bool dflash_laguna = false; + bool dflash_dsv4 = false; // needed by encoder-decoder models (e.g. T5, FLAN-T5) // ref: https://github.com/ggerganov/llama.cpp/pull/8141 @@ -195,6 +196,7 @@ struct llama_hparams { if (this->dflash_n_target_features != other.dflash_n_target_features) return true; if (this->dflash_n_target_layers != other.dflash_n_target_layers) return true; if (this->dflash_laguna != other.dflash_laguna) return true; + if (this->dflash_dsv4 != other.dflash_dsv4) return true; if (this->n_layer != other.n_layer) return true; if (this->n_rot != other.n_rot) return true; if (this->n_swa != other.n_swa) return true; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 4a5c5633d..790927f5d 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -103,6 +103,8 @@ struct create_tensors_helper : public create_tensors_helper_interface { bool create_dflash_tensors(const LLM_TN & tn); + bool create_dflash_dsv4_tensors(const LLM_TN & tn); + bool create_starcoder2_tensors(const LLM_TN & tn); bool create_mamba_tensors(const LLM_TN & tn); @@ -2346,6 +2348,24 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) { model.output_mtp = model.output; model.dflash_fc = create_tensor(ctx_output, tn(LLM_TENSOR_DFLASH_FC, "weight"), {(int64_t) hparams.dflash_n_target_features, n_embd}, 0); model.dflash_hidden_norm = create_tensor(ctx_output, tn(LLM_TENSOR_DFLASH_HIDDEN_NORM, "weight"), {n_embd}, 0); + + const ggml_tensor * markov_w1_meta = ml.get_tensor_meta("markov_w1.weight"); + if (markov_w1_meta != nullptr) { + const int64_t markov_rank = markov_w1_meta->ne[0]; + model.dspark_markov_w1 = create_tensor( + ctx_output, tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), {markov_rank, n_vocab}, + llama_model_loader::TENSOR_NOT_REQUIRED); + model.dspark_markov_w2 = create_tensor( + ctx_output, tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), {markov_rank, n_vocab}, + llama_model_loader::TENSOR_NOT_REQUIRED); + model.dspark_conf_proj = create_tensor( + ctx_output, tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), {n_embd + markov_rank, 1}, + llama_model_loader::TENSOR_NOT_REQUIRED); + model.dspark_conf_proj_b = create_tensor( + ctx_output, tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), {1}, + llama_model_loader::TENSOR_NOT_REQUIRED); + } + model.dflash_aux_hidden_norms.clear(); if (hparams.dflash_laguna) { GGML_ASSERT(hparams.dflash_n_target_layers > 0); @@ -2396,6 +2416,102 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) { return use_mmap_buffer; } +bool create_tensors_helper::create_dflash_dsv4_tensors(const LLM_TN & tn) { + LOADING_PRELUDE + + // DSV4 draft share token embedding and output head. + model.tok_embd = nullptr; + model.output = nullptr; + model.output_mtp = nullptr; + + const std::string markov_w1_name = tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"); + const std::string markov_w2_name = tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"); + const ggml_tensor * markov_w1_meta = ml.require_tensor_meta(markov_w1_name.c_str()); + const ggml_tensor * markov_w2_meta = ml.require_tensor_meta(markov_w2_name.c_str()); + if (ggml_n_dims(markov_w1_meta) != 2 || ggml_n_dims(markov_w2_meta) != 2 || + markov_w1_meta->ne[0] <= 0 || markov_w1_meta->ne[1] != (int64_t) n_vocab || + markov_w2_meta->ne[0] != markov_w1_meta->ne[0] || markov_w2_meta->ne[1] != markov_w1_meta->ne[1]) { + throw std::runtime_error(format("%s: invalid DSV4 Markov tensor shapes: %s=%s, %s=%s", + __func__, markov_w1_name.c_str(), llama_format_tensor_shape(markov_w1_meta).c_str(), + markov_w2_name.c_str(), llama_format_tensor_shape(markov_w2_meta).c_str())); + } + const int64_t markov_rank = markov_w1_meta->ne[0]; + + model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + model.dflash_fc = create_tensor(ctx_output, tn(LLM_TENSOR_DFLASH_FC, "weight"), + {(int64_t) hparams.dflash_n_target_features, n_embd}, 0); + model.dflash_hidden_norm = create_tensor(ctx_output, tn(LLM_TENSOR_DFLASH_HIDDEN_NORM, "weight"), {n_embd}, 0); + model.hc_head_base = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {(int64_t) hparams.dsv4_hc_mult}, 0); + model.hc_head_fn = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_FN, "weight"), + {(int64_t) n_embd * hparams.dsv4_hc_mult, (int64_t) hparams.dsv4_hc_mult}, 0); + model.hc_head_scale = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0); + model.dspark_markov_w1 = create_tensor(ctx_output, markov_w1_name, {markov_rank, n_vocab}, 0); + model.dspark_markov_w2 = create_tensor(ctx_output, markov_w2_name, {markov_rank, n_vocab}, 0); + model.dspark_conf_proj = create_tensor(ctx_output, tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), + {n_embd + markov_rank, 1}, 0); + model.dspark_conf_proj_b = create_tensor(ctx_output, tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), + {1}, llama_model_loader::TENSOR_NOT_REQUIRED); + + for (int i = 0; i < n_layer; ++i) { + ggml_context * ctx_split = ctx_for_layer_split(i); + auto & layer = model.layers[i]; + const int64_t n_head = hparams.n_head(i); + const int64_t n_embd_head = hparams.n_embd_head_k(i); + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t o_group_dim = hparams.dsv4_o_group_count > 0 + ? (n_head * n_embd_head) / hparams.dsv4_o_group_count : 0; + const int64_t o_lora_width = (int64_t) hparams.dsv4_o_lora_rank * hparams.dsv4_o_group_count; + const int64_t hc_width = (int64_t) hparams.dsv4_hc_mult * hparams.dsv4_hc_mult + + 2 * hparams.dsv4_hc_mult; + const std::string ffn_gate_exps_name = tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i); + const ggml_tensor * ffn_gate_exps_meta = ml.require_tensor_meta(ffn_gate_exps_name.c_str()); + if (ggml_n_dims(ffn_gate_exps_meta) != 3 || ffn_gate_exps_meta->ne[0] != n_embd || + ffn_gate_exps_meta->ne[1] <= 0 || ffn_gate_exps_meta->ne[2] != n_expert) { + throw std::runtime_error(format("%s: invalid DSV4 expert tensor shape at layer %d: %s=%s", + __func__, i, ffn_gate_exps_name.c_str(), llama_format_tensor_shape(ffn_gate_exps_meta).c_str())); + } + const int64_t n_ff = ffn_gate_exps_meta->ne[1]; + if (q_lora_rank <= 0 || kv_lora_rank <= 0 || o_group_dim <= 0 || o_lora_width <= 0 || hc_width <= 0) { + throw std::runtime_error(format("%s: invalid DSV4 derived tensor dimensions at layer %d", __func__, i)); + } + + layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_sinks = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0); + layer.wq_a = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0); + layer.attn_q_a_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0); + layer.wq_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0); + layer.wkv_latent = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_KV_LATENT, "weight", i), {n_embd, kv_lora_rank}, 0); + layer.wkv_b = layer.wkv_latent; + layer.wkv_a_mqa = layer.wkv_latent; + layer.attn_kv_a_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0); + layer.attn_kv_norm = layer.attn_kv_a_norm; + layer.wo_a = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {o_group_dim, o_lora_width}, 0); + layer.wo_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_lora_width, n_embd}, 0); + layer.wo = layer.wo_b; + + layer.hc_attn_base = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_width}, 0); + layer.hc_attn_fn = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {n_embd * hparams.dsv4_hc_mult, hc_width}, 0); + layer.hc_attn_scale = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0); + layer.hc_ffn_base = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_width}, 0); + layer.hc_ffn_fn = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {n_embd * hparams.dsv4_hc_mult, hc_width}, 0); + layer.hc_ffn_scale = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0); + + layer.ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_gate_inp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_gate_exps = create_tensor(ctx_split, ffn_gate_exps_name, {n_embd, n_ff, n_expert}, 0); + layer.ffn_down_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0); + const int64_t n_ff_shexp = n_ff * hparams.n_expert_shared; + layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0); + layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_exp_probs_b = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + } + + return use_mmap_buffer; +} + bool create_tensors_helper::create_starcoder2_tensors(const LLM_TN & tn) { LOADING_PRELUDE @@ -4932,6 +5048,8 @@ bool create_tensors_helper::create_tensors() { case LLM_ARCH_GEMMA4_MTP: case LLM_ARCH_GEMMA4_ASSISTANT: use_mmap_buffer = create_gemma4_mtp_tensors(tn); break; + case LLM_ARCH_DFLASH: + use_mmap_buffer = create_dflash_dsv4_tensors(tn); break; case LLM_ARCH_DFLASH_DRAFT: use_mmap_buffer = create_dflash_tensors(tn); break; case LLM_ARCH_STARCODER2: diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 5a24fdd16..f347c744d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -868,6 +868,49 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_DFLASH_FC, "dflash_fc" }, { LLM_TENSOR_DFLASH_HIDDEN_NORM, "dflash_hidden_norm" }, { LLM_TENSOR_DFLASH_AUX_HIDDEN_NORM, "dflash_aux_hidden_norm.%d" }, + { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, + { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, + { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, + }, + }, + { + LLM_ARCH_DFLASH, + { + { LLM_TENSOR_TOKEN_EMBD, "token_embd" }, + { LLM_TENSOR_OUTPUT_NORM, "output_norm" }, + { LLM_TENSOR_OUTPUT, "output" }, + { LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" }, + { LLM_TENSOR_ATTN_SINKS, "blk.%d.attn_sinks" }, + { LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" }, + { LLM_TENSOR_ATTN_KV_A_NORM, "blk.%d.attn_kv_a_norm" }, + { LLM_TENSOR_ATTN_Q_A, "blk.%d.attn_q_a" }, + { LLM_TENSOR_ATTN_Q_B, "blk.%d.attn_q_b" }, + { LLM_TENSOR_ATTN_KV_LATENT, "blk.%d.attn_kv" }, + { LLM_TENSOR_ATTN_OUT_A, "blk.%d.attn_output_a" }, + { LLM_TENSOR_ATTN_OUT_B, "blk.%d.attn_output_b" }, + { LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" }, + { LLM_TENSOR_FFN_GATE_INP, "blk.%d.ffn_gate_inp" }, + { LLM_TENSOR_FFN_GATE_EXPS, "blk.%d.ffn_gate_exps" }, + { LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" }, + { LLM_TENSOR_FFN_UP_EXPS, "blk.%d.ffn_up_exps" }, + { LLM_TENSOR_FFN_GATE_SHEXP, "blk.%d.ffn_gate_shexp" }, + { LLM_TENSOR_FFN_DOWN_SHEXP, "blk.%d.ffn_down_shexp" }, + { LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" }, + { LLM_TENSOR_FFN_EXP_PROBS_B, "blk.%d.exp_probs_b" }, + { LLM_TENSOR_HC_HEAD_BASE, "output_hc_base" }, + { LLM_TENSOR_HC_HEAD_FN, "output_hc_fn" }, + { LLM_TENSOR_HC_HEAD_SCALE, "output_hc_scale" }, + { LLM_TENSOR_HC_ATTN_BASE, "blk.%d.hc_attn_base" }, + { LLM_TENSOR_HC_ATTN_FN, "blk.%d.hc_attn_fn" }, + { LLM_TENSOR_HC_ATTN_SCALE, "blk.%d.hc_attn_scale" }, + { LLM_TENSOR_HC_FFN_BASE, "blk.%d.hc_ffn_base" }, + { LLM_TENSOR_HC_FFN_FN, "blk.%d.hc_ffn_fn" }, + { LLM_TENSOR_HC_FFN_SCALE, "blk.%d.hc_ffn_scale" }, + { LLM_TENSOR_DFLASH_FC, "fc" }, + { LLM_TENSOR_DFLASH_HIDDEN_NORM, "enc.output_norm" }, + { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, + { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, + { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, }, }, { diff --git a/src/llama-model.h b/src/llama-model.h index 71c51a923..891a3772b 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -474,6 +474,10 @@ struct llama_model { struct ggml_tensor * dflash_fc = nullptr; struct ggml_tensor * dflash_hidden_norm = nullptr; std::vector dflash_aux_hidden_norms; + struct ggml_tensor * dspark_markov_w1 = nullptr; + struct ggml_tensor * dspark_markov_w2 = nullptr; + struct ggml_tensor * dspark_conf_proj = nullptr; + struct ggml_tensor * dspark_conf_proj_b = nullptr; struct ggml_tensor * output_norm; struct ggml_tensor * output_norm_b; diff --git a/src/llama-spec-features-dflash.cpp b/src/llama-spec-features-dflash.cpp index 7348388a6..deed94a22 100644 --- a/src/llama-spec-features-dflash.cpp +++ b/src/llama-spec-features-dflash.cpp @@ -75,6 +75,12 @@ int32_t llama_get_dflash_visible_cross_ctx( return ctx != nullptr ? ctx->dflash.visible_cross_ctx : 0; } +void llama_set_dflash_dspark(struct llama_context * ctx, bool enabled) { + if (ctx != nullptr) { + ctx->dflash.dspark = enabled; + } +} + int32_t llama_model_dflash_block_size(const struct llama_model * model) { return model ? (int32_t) model->hparams.dflash_block_size : 0; } @@ -115,6 +121,12 @@ int32_t llama_model_dflash_target_mask_token_id(const struct llama_model * model return (int32_t) model->vocab.token_mask(); } +bool llama_model_dflash_has_dspark_head(const struct llama_model * model) { + return model != nullptr && + model->dspark_markov_w1 != nullptr && + model->dspark_markov_w2 != nullptr; +} + static const ggml_tensor * llama_dflash_output_tensor( const struct llama_model * model) { if (model == nullptr) { @@ -135,7 +147,7 @@ static const ggml_tensor * llama_dflash_output_tensor( int32_t llama_model_dflash_io_mode( const struct llama_model * draft_model, const struct llama_model * target_model) { - if (draft_model == nullptr || target_model == nullptr || draft_model->arch != LLM_ARCH_DFLASH_DRAFT) { + if (draft_model == nullptr || target_model == nullptr || !llm_arch_is_dflash_family(draft_model->arch)) { return LLAMA_DFLASH_IO_MODE_INVALID; } @@ -219,7 +231,7 @@ bool llama_model_share_dflash_io_tensors( return false; } - if (draft_model->arch != LLM_ARCH_DFLASH_DRAFT) { + if (!llm_arch_is_dflash_family(draft_model->arch)) { return true; } diff --git a/src/llama-spec-features-dflash.h b/src/llama-spec-features-dflash.h index f0a1fde6c..7e11d4fdf 100644 --- a/src/llama-spec-features-dflash.h +++ b/src/llama-spec-features-dflash.h @@ -78,6 +78,7 @@ llama_dflash_kv_cache_transition llama_plan_dflash_kv_cache_transition_for_ctx( void llama_reset_dflash_kv_cache_state(struct llama_context * ctx); void llama_set_dflash_visible_cross_ctx(struct llama_context * ctx, int32_t cross_ctx); int32_t llama_get_dflash_visible_cross_ctx(const struct llama_context * ctx); +void llama_set_dflash_dspark(struct llama_context * ctx, bool enabled); int32_t llama_model_dflash_block_size(const struct llama_model * model); int32_t llama_model_dflash_mask_token_id(const struct llama_model * model); @@ -85,6 +86,7 @@ int32_t llama_model_dflash_n_target_layers(const struct llama_model * model); int32_t llama_model_dflash_n_target_features(const struct llama_model * model); int32_t llama_model_dflash_target_layer_ids(const struct llama_model * model, int32_t * layer_ids, int32_t capacity); int32_t llama_model_dflash_target_mask_token_id(const struct llama_model * model); +bool llama_model_dflash_has_dspark_head(const struct llama_model * model); enum llama_dflash_io_mode { LLAMA_DFLASH_IO_MODE_INVALID = 0, diff --git a/src/llama.cpp b/src/llama.cpp index 292ef9739..1bf9f75ff 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -3970,7 +3970,11 @@ static std::pair, double> get_layer_sizes(const llama_model_ continue; } if (name == "dflash_fc.weight" || name == "dflash_hidden_norm.weight" || - name.rfind("dflash_aux_hidden_norm.", 0) == 0) { + (model.arch == LLM_ARCH_DFLASH && + (name == "fc.weight" || name == "enc.output_norm.weight")) || + name.rfind("dflash_aux_hidden_norm.", 0) == 0 || + name == "markov_w1.weight" || name == "markov_w2.weight" || + name == "conf_proj.weight" || name == "conf_proj.bias") { output_misc_size += size; continue; } @@ -4824,7 +4828,7 @@ static bool llm_load_tensors( if (model.arch == LLM_ARCH_GEMMA4) { llm_scale_gate_inp_s(model, use_mmap_buffer); } - if ((model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DFLASH_DRAFT) && extra_output_type != GGML_TYPE_COUNT) { + if ((model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || llm_arch_is_dflash_family(model.arch)) && extra_output_type != GGML_TYPE_COUNT) { llm_requantize_output_tensor(model, extra_output_type); } @@ -6159,7 +6163,7 @@ static int llama_decode_internal( // reserve output buffer n_outputs_embd = has_mtp && cparams.mtp_op_type == MTP_OP_NONE ? n_tokens_all : n_outputs; const size_t required_outputs = std::max(n_outputs, n_outputs_embd); - const bool is_dflash_decode = lctx.model.arch == LLM_ARCH_DFLASH_DRAFT; + const bool is_dflash_decode = llm_arch_is_dflash_family(lctx.model.arch); const size_t reserved_outputs = llama_output_reserve(lctx, required_outputs); if (reserved_outputs < required_outputs) { LLAMA_LOG_ERROR("%s: could not reserve space for batch with %zu outputs\n", __func__, required_outputs); @@ -6560,7 +6564,7 @@ static int llama_decode_internal( // extract logits { - const bool dflash_skip_logits = (lctx.model.arch == LLM_ARCH_DFLASH_DRAFT + const bool dflash_skip_logits = (llm_arch_is_dflash_family(lctx.model.arch) && !lctx.dflash.draft_tokens.empty()); if (dflash_skip_logits) { res = nullptr; @@ -8820,6 +8824,8 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) { case LLM_ARCH_LAGUNA: case LLM_ARCH_GEMMA4: case LLM_ARCH_GEMMA4_MTP: + case LLM_ARCH_DFLASH: + return LLAMA_ROPE_TYPE_NORM; case LLM_ARCH_DFLASH_DRAFT: case LLM_ARCH_GEMMA4_ASSISTANT: return LLAMA_ROPE_TYPE_NEOX;