Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-01 10:00:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

DeepseekV4 MTP + DSpark (#25784)

Browse Source
This commit is contained in:
Aman Gupta
2026-08-02 20:55:34 +08:00
committed by GitHub
parent f5919bf458
commit 596a5795bd
14 changed files with 1534 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/dflash.cpp
+270
View File
@@ -20,6 +20,48 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
}
LLAMA_LOG_INFO("]\n");
// DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);
if (hparams.dsv4_hc_mult > 0) {
ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale);
ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm);
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all);
if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {
hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;
}
ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count);
ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank);
ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false);
if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {
throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring");
}
for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
if (hparams.dsv4_compress_ratios[il] != 0) {
throw std::runtime_error("DSpark DSV4 draft expects uncompressed attention on all stages");
}
}
GGML_ASSERT(hparams.n_swa > 0);
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
hparams.set_swa_pattern(0);
for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
hparams.is_swa_impl[il] = true;
}
hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
type = LLM_TYPE_UNKNOWN;
return;
}
// optional interleaved sliding-window attention with per-layer pattern array.
// DFlash has a single rope, so the SWA rope == main rope.
if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
@@ -58,6 +100,56 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm
if (hparams.dsv4_hc_mult > 0) {
const int64_t q_lora_rank = hparams.n_lora_q;
const int64_t n_ff_exp = hparams.n_ff_exp;
const int64_t n_expert_shared = hparams.n_expert_shared;
const int64_t n_embd_head = hparams.n_embd_head_k();
const int64_t o_groups = hparams.dsv4_o_group_count;
const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;
const int64_t hc_mult = hparams.dsv4_hc_mult;
const int64_t hc_dim = hc_mult * n_embd;
const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0);
hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0);
hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);
layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);
layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);
layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);
layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);
layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, 0);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);
layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);
layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);
layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);
layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
}
return;
}
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
@@ -84,6 +176,9 @@ std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const ll
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
if (hparams.dsv4_hc_mult > 0) {
return std::make_unique<graph_dsv4>(*this, params);
}
return std::make_unique<graph<false>>(*this, params);
default:
GGML_ABORT("invalid graph type");
@@ -403,3 +498,178 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
build_dspark_markov_head(*this, model, inp_tokens);
}
}
// DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above):
// * embd batch -> project main_x through each stage's wkv and inject K into the ring cache
// * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads
llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_graph_params & params) :
llama_model_deepseek4::graph(params) {
const int64_t n_embd_head = hparams.n_embd_head_k();
const int64_t n_embd_head_rope = hparams.n_rot();
const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;
ggml_tensor * inp_pos = build_inp_pos();
llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
// KV cache injection: fused target features from the encoder
if (ubatch.embd) {
auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
ggml_set_input(inp->embd);
ggml_tensor * inp_g = inp->embd;
cb(inp_g, "inp_g_embeddings", -1);
res->add_input(std::move(inp));
for (int il = 0; il < n_layer; ++il) {
const auto & layer = model.layers[il];
// main-track KV: kv_norm(wkv(main_x)) with rope on the trailing dims, same
// rope parameters as the uncompressed layers in build_attention_impl
ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);
kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);
kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);
ggml_tensor * kv_nope = ggml_view_3d(ctx0, kv, n_embd_head_nope, 1, n_tokens,
ggml_row_size(kv->type, n_embd_head),
ggml_row_size(kv->type, n_embd_head),
0);
ggml_tensor * kv_pe = ggml_view_3d(ctx0, kv, n_embd_head_rope, 1, n_tokens,
ggml_row_size(kv->type, n_embd_head),
ggml_row_size(kv->type, n_embd_head),
ggml_row_size(kv->type, n_embd_head_nope));
kv_pe = ggml_rope_ext(ctx0, kv_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,
freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
kv = ggml_concat(ctx0, kv_nope, kv_pe, 0);
cb(kv, "kv_injected", il);
if (inp_attn->self_k_rot_swa) {
kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);
}
ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
}
res->t_embd = inp_g;
ggml_build_forward_expand(gf, inp_g);
return;
}
// tok_embd from the target model (shared via ctx_other)
auto * tok_embd = model.tok_embd;
if (tok_embd == nullptr) {
GGML_ASSERT(cparams.ctx_other != nullptr);
const auto * model_other = llama_get_model(cparams.ctx_other);
GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");
tok_embd = model_other->tok_embd;
}
auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
ggml_tensor * inp_tokens = inp->tokens;
ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
cb(inpL, "inp_noise_embd", -1);
res->add_input(std::move(inp));
const int64_t hc = hparams.dsv4_hc_mult;
inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);
inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
cb(inpL, "hc_init", -1);
for (int il = 0; il < n_layer; ++il) {
const auto & layer = model.layers[il];
ggml_tensor * residual = inpL;
ggml_tensor * post = nullptr;
ggml_tensor * comb = nullptr;
ggml_tensor * cur = build_hc_pre(inpL,
layer.hc_attn_fn,
layer.hc_attn_scale,
layer.hc_attn_base,
&post, &comb, il);
cb(cur, "hc_attn_pre", il);
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "attn_norm", il);
cur = build_attention(model, inp_attn, cur, inp_pos, il);
inpL = build_hc_post(cur, residual, post, comb, il);
cb(inpL, "hc_attn_post", il);
residual = inpL;
cur = build_hc_pre(inpL,
layer.hc_ffn_fn,
layer.hc_ffn_scale,
layer.hc_ffn_base,
&post, &comb, il);
cb(cur, "hc_ffn_pre", il);
cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "ffn_norm", il);
ggml_tensor * moe_out = build_moe_ffn(cur,
layer.ffn_gate_inp,
layer.ffn_up_exps,
layer.ffn_gate_exps,
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
n_expert, hparams.n_expert_used,
LLM_FFN_SILU, hparams.expert_weights_norm,
hparams.expert_weights_scale,
(llama_expert_gating_func_type) hparams.expert_gating_func,
il);
cb(moe_out, "ffn_moe_out", il);
ggml_tensor * ffn_shexp = build_ffn(cur,
layer.ffn_up_shexp, nullptr, nullptr,
layer.ffn_gate_shexp, nullptr, nullptr,
layer.ffn_down_shexp, nullptr, nullptr,
nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(ffn_shexp, "ffn_shexp", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
cb(cur, "ffn_out", il);
inpL = build_hc_post(cur, residual, post, comb, il);
cb(inpL, "l_out", il);
}
ggml_tensor * cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
cb(cur, "hc_head", -1);
// confidence head input: the reference scores the pre-norm collapsed hidden state
res->t_embd = cur;
cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
// lm_head from the target model (shared via ctx_other)
auto * output = model.output;
if (output == nullptr) {
GGML_ASSERT(cparams.ctx_other != nullptr);
const auto * model_other = llama_get_model(cparams.ctx_other);
GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
output = model_other->output;
}
cur = build_lora_mm(output, cur);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
if (model.dspark_markov_w1) {
build_dspark_markov_head(*this, model, inp_tokens);
}
}