LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 14:24:03 +04:00
Code Issues Packages Projects Releases Wiki Activity

Initial implementation of DSpark (#2280)

Browse Source 
* Implement initial arch for DSpark

* feat: Add Dspark architecture support

* avoid to many splits in graph and improve rope logic
This commit is contained in:
Samuel Oliveira Alves
2026-08-10 08:46:03 +02:00
committed by GitHub
parent a7c81affa4
commit 7ebbb906d2
25 changed files with 776 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+27 -5
View File
@@ -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"
common/common.h
+5
View File
@@ -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;
common/speculative-dflash-impl.h
+24 -5
View File
@@ -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<int32_t> 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<int32_t>(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<int32_t>(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);
}
common/speculative.cpp
+34 -13
View File
@@ -30,6 +30,7 @@ const std::vector<enum common_speculative_type> 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<std::string, enum common_speculative_type> 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(), [&params](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<common_speculative_state_dflash>(
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<spec_tuner>();
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<int32_t> & output_indices,
std::vector<float> & 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;
}
convert_hf_to_gguf.py
+60
View File
@@ -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
examples/main/main.cpp
+1 -2
View File
@@ -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) {
examples/server/server-context.cpp
+1 -2
View File
@@ -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) {
gguf-py/gguf/constants.py
+10
View File
@@ -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,
src/graphs/build_deepseek4.cpp
+194 -1
View File
@@ -8,6 +8,17 @@
#include <cstring>
#include <string>
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<int64_t>(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<int64_t>(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;
}
src/graphs/build_dflash.cpp
+110 -2
View File
@@ -4,6 +4,53 @@
#include <cmath>
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;
src/llama-arch.cpp
+6
View File
@@ -85,6 +85,7 @@ static const std::map<llm_arch, const char *> 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, const char *> 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;
}
src/llama-arch.h
+6
View File
@@ -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);
src/llama-build-context.cpp
+1
View File
@@ -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();
src/llama-build-context.h
+6
View File
@@ -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,
src/llama-context.h
+2 -2
View File
@@ -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<float> 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<llama_token> draft_tokens;
struct ggml_tensor * draft_tokens_tensor = nullptr;
};
src/llama-dflash.cpp
+39 -7
View File
@@ -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<int32_t>(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<int32_t> 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;
}
}
src/llama-dsv4.cpp
+18 -5
View File
@@ -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);
}
}
src/llama-hparams.cpp
+39 -6
View File
@@ -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;
}
src/llama-hparams.h
+2
View File
@@ -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;
src/llama-load-tensors.cpp
+118
View File
@@ -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:
src/llama-model.cpp
+43
View File
@@ -868,6 +868,49 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> 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" },
},
},
{
src/llama-model.h
+4
View File
@@ -474,6 +474,10 @@ struct llama_model {
struct ggml_tensor * dflash_fc = nullptr;
struct ggml_tensor * dflash_hidden_norm = nullptr;
std::vector<struct ggml_tensor *> 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;
src/llama-spec-features-dflash.cpp
+14 -2
View File
@@ -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;
}
src/llama-spec-features-dflash.h
+2
View File
@@ -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,
src/llama.cpp
+10 -4
View File
@@ -3970,7 +3970,11 @@ static std::pair<std::vector<double>, 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<size_t>(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;