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

Merge commit '0b14b87d7c20cb753b94b96854dd7b45306fc696' into concedo_experimental

Browse Source 
# Conflicts:
#	.github/workflows/build-cpu.yml
#	.github/workflows/release.yml
#	AGENTS.md
#	ggml/src/ggml-opencl/ggml-opencl.cpp
#	ggml/src/ggml-sycl/common.hpp
#	ggml/src/ggml-sycl/ggml-sycl.cpp
#	ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp
#	ggml/src/ggml-webgpu/ggml-webgpu.cpp
#	ggml/src/ggml-webgpu/wgsl-shaders/repeat.wgsl
#	src/models/qwen3next.cpp
#	tests/test-backend-ops.cpp
#	tests/test-chat-peg-parser.cpp
#	tests/test-chat.cpp
#	tests/test-llama-archs.cpp
#	tests/test-recurrent-state-rollback.cpp
This commit is contained in:
Concedo
2026-08-07 18:04:33 +08:00
parent 569532f7e4 0b14b87d7c
commit 7098c9c53f
67 changed files with 4134 additions and 807 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-arch.cpp
+2
View File
@@ -968,6 +968,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_DEEPSEEK4:
return true;
default:
return false;
@@ -990,6 +991,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
switch (arch) {
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_DEEPSEEK4:
return true;
default:
return false;
src/llama-context.cpp
+14 -10
View File
@@ -123,8 +123,9 @@ llama_context::llama_context(
cparams.no_perf = params.no_perf;
cparams.warmup = false;
cparams.embeddings_layer_inp.resize(hparams.n_layer(), false);
embd_layer_inp.resize(hparams.n_layer());
// +1: id n_layer() taps the output of the last layer ("input" of the head)
cparams.embeddings_layer_inp.resize(hparams.n_layer() + 1, false);
embd_layer_inp.resize(hparams.n_layer() + 1);
cparams.ctx_type = params.ctx_type;
cparams.pooling_type = params.pooling_type;
@@ -1174,7 +1175,7 @@ void llama_context::set_embeddings_nextn(bool value, bool masked) {
void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) {
LLAMA_LOG_DEBUG("%s: lid = %d, enable = %d\n", __func__, lid, enable);
GGML_ASSERT(lid < model.hparams.n_layer());
GGML_ASSERT(lid <= model.hparams.n_layer());
cparams.embeddings_layer_inp[lid] = enable;
@@ -1726,7 +1727,8 @@ int llama_context::decode(const llama_batch & batch_inp) {
const auto & hparams = model.hparams;
const int64_t n_vocab = vocab.n_tokens();
const int64_t n_embd = hparams.n_embd_inp();
const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd;
const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : hparams.n_embd_inp();
// when computing embeddings, all tokens are output
const bool output_all = cparams.embeddings;
@@ -2285,8 +2287,9 @@ void llama_context::extract_layer_inputs(const llm_graph_result * res, size_t to
}
void llama_context::output_reorder() {
const uint64_t n_vocab = model.vocab.n_tokens();
const uint64_t n_embd = model.hparams.n_embd;
const uint64_t n_vocab = model.vocab.n_tokens();
const uint64_t n_embd = model.hparams.n_embd;
const uint64_t n_embd_out = model.hparams.n_embd_out();
for (size_t s = 0; s < output_swaps.size(); ++s) {
const uint64_t i0 = output_swaps[s].i0;
@@ -2299,14 +2302,14 @@ void llama_context::output_reorder() {
}
if (embd.size > 0) {
for (uint64_t k = 0; k < n_embd; k++) {
std::swap(embd.data[i0*n_embd + k], embd.data[i1*n_embd + k]);
for (uint64_t k = 0; k < n_embd_out; k++) {
std::swap(embd.data[i0*n_embd_out + k], embd.data[i1*n_embd_out + k]);
}
}
if (embd_nextn.size > 0) {
for (uint64_t k = 0; k < n_embd; k++) {
std::swap(embd_nextn.data[i0*n_embd + k], embd_nextn.data[i1*n_embd + k]);
for (uint64_t k = 0; k < n_embd_out; k++) {
std::swap(embd_nextn.data[i0*n_embd_out + k], embd_nextn.data[i1*n_embd_out + k]);
}
}
@@ -2361,6 +2364,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE ||
model.arch == LLM_ARCH_DEEPSEEK4 ||
(model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
model.arch == LLM_ARCH_NANBEIGE ||
model.arch == LLM_ARCH_MINIMAX_M3) {
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
src/llama-graph.cpp
+169 -3
View File
@@ -619,6 +619,63 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
return res;
}
void llm_graph_input_attn_k_iswa::set_input(const llama_ubatch * ubatch) {
// base tensors may not be allocated if there are no non-SWA attention layers
if (self_k_idxs && self_k_idxs->buffer) {
mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch);
}
// the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live
if (self_kq_mask && self_kq_mask->buffer) {
mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
}
// swa tensors may not be allocated if there are no SWA attention layers
if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch);
}
if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
}
if (self_k_rot && self_k_rot->buffer) {
mctx->get_base()->set_input_k_rot(self_k_rot);
}
if (self_k_rot_swa && self_k_rot_swa->buffer) {
mctx->get_swa()->set_input_k_rot(self_k_rot_swa);
}
}
bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) {
const auto * mctx = static_cast<const llama_kv_cache_iswa_context *>(params.mctx);
this->mctx = mctx;
bool res = true;
// base tensors may not be allocated if there are no non-SWA attention layers
if (self_k_idxs && self_k_idxs->buffer) {
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
}
if (self_kq_mask && self_kq_mask->buffer) {
res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams);
}
// swa tensors may not be allocated if there are no SWA attention layers
if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;
}
if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
}
return res;
}
static void dsv4_set_i64(ggml_tensor * dst, const std::vector<int64_t> & src) {
if (!dst || !dst->buffer) {
return;
@@ -755,6 +812,10 @@ static void dsv4_set_comp_inputs(
dsv4_set_i32(inp.state_pos, plan.state_pos);
dsv4_set_i32(inp.state_persist_src_idxs, plan.state_persist_src_idxs);
dsv4_set_i32(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs);
dsv4_set_i32(inp.state_restore_src_idxs, plan.state_restore_src_idxs);
dsv4_set_i32(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs);
dsv4_set_i32(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs);
dsv4_set_i32(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs);
dsv4_set_i32(inp.state_read_idxs, plan.state_read_idxs);
dsv4_set_i64(inp.state_write_idxs, plan.state_write_idxs);
dsv4_set_i32(inp.state_write_pos, plan.state_write_pos);
@@ -799,6 +860,10 @@ static bool dsv4_can_reuse_comp_input(
res &= dsv4_can_reuse_tensor_1d(inp.state_pos, plan.state_pos.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_persist_src_idxs, plan.state_persist_src_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_restore_src_idxs, plan.state_restore_src_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_read_idxs, plan.state_read_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_write_idxs, plan.state_write_idxs.size());
res &= dsv4_can_reuse_tensor_1d(inp.state_write_pos, plan.state_write_pos.size());
@@ -833,6 +898,10 @@ static void dsv4_build_comp_inputs(
inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos");
inp.state_persist_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_src_idxs.size(), std::string("dsv4_") + name + "_state_persist_src_idxs");
inp.state_persist_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_dst_idxs.size(), std::string("dsv4_") + name + "_state_persist_dst_idxs");
inp.state_restore_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_src_idxs.size(), std::string("dsv4_") + name + "_state_restore_src_idxs");
inp.state_restore_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_dst_idxs.size(), std::string("dsv4_") + name + "_state_restore_dst_idxs");
inp.state_snapshot_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_src_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_src_idxs");
inp.state_snapshot_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_dst_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_dst_idxs");
inp.state_read_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_read_idxs.size(), std::string("dsv4_") + name + "_state_read_idxs");
inp.state_write_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I64, plan.state_write_idxs.size(), std::string("dsv4_") + name + "_state_write_idxs");
inp.state_write_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_write_pos.size(), std::string("dsv4_") + name + "_state_write_pos");
@@ -1196,7 +1265,7 @@ void llm_graph_result::reset() {
t_embd_pooled = nullptr;
t_h_nextn = nullptr;
t_layer_inp.resize(LLAMA_MAX_LAYERS);
t_layer_inp.resize(LLAMA_MAX_LAYERS + 1);
std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr);
t_sampled.clear();
@@ -1651,7 +1720,7 @@ ggml_tensor * llm_graph_context::build_ffn(
tmp = ggml_clamp(ctx0, tmp, -limit, limit);
cb(tmp, "ffn_up_clamped", il);
if (arch == LLM_ARCH_DEEPSEEK4) {
if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
cur = ggml_clamp(ctx0, cur, -INFINITY, limit);
cb(cur, "ffn_gate_clamped", il);
cur = ggml_swiglu_split(ctx0, cur, tmp);
@@ -2046,7 +2115,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
up = ggml_clamp(ctx0, up, -limit, limit);
cb(up, "ffn_moe_up_clamped", il);
if (arch == LLM_ARCH_DEEPSEEK4) {
if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
cur = ggml_clamp(ctx0, cur, -INFINITY, limit);
cb(cur, "ffn_moe_gate_clamped", il);
cur = ggml_swiglu_split(ctx0, cur, up);
@@ -2963,6 +3032,75 @@ ggml_tensor * llm_graph_context::build_attn(
return cur;
}
ggml_tensor * llm_graph_context::build_attn(
llm_graph_input_attn_k_iswa * inp,
ggml_tensor * wo,
ggml_tensor * wo_b,
ggml_tensor * wo_s,
ggml_tensor * q_cur,
ggml_tensor * k_cur,
ggml_tensor * v_cur,
ggml_tensor * kq_b,
ggml_tensor * sinks,
ggml_tensor * v_mla,
float kq_scale,
int il) const {
const bool is_swa = hparams.is_swa(il);
GGML_UNUSED(v_cur);
auto * k_rot = is_swa ? inp->self_k_rot_swa : inp->self_k_rot;
if (k_rot) {
q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot);
if (k_cur) {
k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot);
}
}
// these nodes are added to the graph together so that they are not reordered
// by doing so, the number of splits in the graph is reduced
ggml_build_forward_expand(gf, q_cur);
if (k_cur) {
ggml_build_forward_expand(gf, k_cur);
}
const auto * mctx_iswa = inp->mctx;
const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
// optionally store to KV cache
if (k_cur) {
const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
}
const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
// MLA-style attention: the cached K is used as V
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = k;
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
cb(cur, "kqv_out", il);
if (k_rot) {
cur = llama_mul_mat_hadamard(ctx0, cur, k_rot);
}
if (wo) {
cur = build_lora_mm(wo, cur, wo_s);
}
if (wo_b) {
cur = ggml_add(ctx0, cur, wo_b);
}
return cur;
}
llm_graph_input_attn_cross * llm_graph_context::build_attn_inp_cross() const {
auto inp = std::make_unique<llm_graph_input_attn_cross>(cross);
@@ -3085,6 +3223,34 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const
return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp));
}
llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const {
const auto * mctx_cur = static_cast<const llama_kv_cache_iswa_context *>(mctx);
auto inp = std::make_unique<llm_graph_input_attn_k_iswa>(hparams, cparams, mctx_cur);
{
inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch);
inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams);
inp->self_kq_mask_cnv = inp->self_kq_mask;
}
{
GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA");
inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);
inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
}
inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0);
inp->self_k_rot_swa = mctx_cur->get_swa()->build_input_k_rot(ctx0);
return (llm_graph_input_attn_k_iswa *) res->add_input(std::move(inp));
}
llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
const auto * mctx_cur = static_cast<const llama_kv_cache_dsv4_context *>(mctx);
const auto * raw_ctx = mctx_cur->get_raw();
src/llama-graph.h
+62 -1
View File
@@ -471,6 +471,45 @@ public:
const llama_kv_cache_iswa_context * mctx;
};
class llm_graph_input_attn_k_iswa : public llm_graph_input_i {
public:
llm_graph_input_attn_k_iswa(
const llama_hparams & hparams,
const llama_cparams & cparams,
const llama_kv_cache_iswa_context * mctx) :
hparams(hparams),
cparams(cparams),
mctx(mctx) {
}
~llm_graph_input_attn_k_iswa() = default;
void set_input(const llama_ubatch * ubatch) override;
bool can_reuse(const llm_graph_params & params) override;
ggml_tensor * get_k_idxs() const { return self_k_idxs; }
ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }
ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]
ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
ggml_tensor * self_k_rot = nullptr;
ggml_tensor * self_k_rot_swa = nullptr;
const llama_hparams hparams;
const llama_cparams cparams;
const llama_kv_cache_iswa_context * mctx;
};
// DSV4 raw graph inputs are SWA-only, but their mask may be stream-shaped
// so raw K can be concatenated with DSV4 compressed K in one attention op.
class llm_graph_input_dsv4_raw {
@@ -505,6 +544,10 @@ public:
ggml_tensor * state_pos = nullptr; // I32 [n_state]
ggml_tensor * state_persist_src_idxs = nullptr; // I32 [n_state_persist]
ggml_tensor * state_persist_dst_idxs = nullptr; // I32 [n_state_persist]
ggml_tensor * state_restore_src_idxs = nullptr; // I32 [n_state_restore]
ggml_tensor * state_restore_dst_idxs = nullptr; // I32 [n_state_restore]
ggml_tensor * state_snapshot_src_idxs = nullptr; // I32 [n_state_snapshot]
ggml_tensor * state_snapshot_dst_idxs = nullptr; // I32 [n_state_snapshot]
ggml_tensor * state_read_idxs = nullptr; // I32 [ratio*n_state_write]
ggml_tensor * state_write_idxs = nullptr; // I64 [n_state_write]
ggml_tensor * state_write_pos = nullptr; // I32 [n_state_write]
@@ -1068,7 +1111,7 @@ struct llm_graph_context {
ggml_tensor * build_attn_mha(
ggml_tensor * q, // [n_embd_head_q, n_head_q, n_tokens]
ggml_tensor * k, // [n_embd_head_k, n_head_k, n_tokens]
ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans == false)
ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans = false)
ggml_tensor * kq_b,
ggml_tensor * kq_mask,
ggml_tensor * sinks, // [n_head_q]
@@ -1160,6 +1203,24 @@ struct llm_graph_context {
float kq_scale,
int il) const;
llm_graph_input_attn_k_iswa * build_attn_inp_k_iswa() const;
// note: if k_cur is not provided, it will not be stored in the memory
// note: the K cache is used as V (MLA-style attention)
ggml_tensor * build_attn(
llm_graph_input_attn_k_iswa * inp,
ggml_tensor * wo,
ggml_tensor * wo_b,
ggml_tensor * wo_s,
ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens]
ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] optional
ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] optional
ggml_tensor * kq_b,
ggml_tensor * sinks, // [n_head_q]
ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]
float kq_scale,
int il) const;
llm_graph_input_attn_cross * build_attn_inp_cross() const;
ggml_tensor * build_attn(
src/llama-kv-cache-dsv4.cpp
+282 -65
View File
@@ -252,7 +252,8 @@ static void dsv4_state_write_tensor_streams(
uint32_t tensor_rows,
uint32_t n_rows,
uint32_t s0,
uint32_t ns) {
uint32_t ns,
const std::vector<uint32_t> * stream_ids = nullptr) {
const int32_t type_i = (int32_t) tensor->type;
const uint64_t ne0 = tensor->ne[0];
const uint64_t rows = n_rows;
@@ -273,8 +274,16 @@ static void dsv4_state_write_tensor_streams(
return;
}
if (stream_ids && stream_ids->size() != ns) {
throw std::runtime_error("DSV4 state tensor stream map size mismatch");
}
for (uint32_t s = 0; s < ns; ++s) {
const size_t offset = (size_t) (s0 + s)*stream_stride;
const uint32_t stream = stream_ids ? (*stream_ids)[s] : s0 + s;
if ((int64_t) stream >= tensor->ne[2]) {
throw std::runtime_error("DSV4 state tensor stream out of range");
}
const size_t offset = (size_t) stream*stream_stride;
io.write_tensor(tensor, offset, size);
}
}
@@ -421,7 +430,9 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
bool overlap,
uint32_t state_size,
uint32_t kv_size,
uint32_t n_stream) {
uint32_t n_stream,
uint32_t n_rs_seq,
const std::vector<uint32_t> & rs_idx) {
llama_kv_cache_dsv4_context::comp_plan plan;
plan.n_visible.resize(ubatch.n_tokens);
plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);
@@ -451,6 +462,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
std::vector<int32_t> overlap_cur_reads;
std::map<std::pair<llama_seq_id, llama_pos>, int64_t> curr_token_idx_map;
std::map<llama_seq_id, uint32_t> state_write_counts;
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
@@ -513,6 +525,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
plan.state_write_idxs.push_back(cache_off + pos/ratio);
plan.state_write_pos.push_back((int32_t) source_start);
++state_write_counts[seq_id];
if (overlap) {
const llama_pos prev_start = source_start - ratio;
@@ -531,33 +544,57 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
}
}
if (ratio == DSV4_CSA_RATIO && plan.state_write_idxs.empty() && !plan.state_pos.empty()) {
// Non-boundary CSA steps still need a write op so their graph matches
// boundary steps. Use a padded scratch row that is masked from attention.
if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) {
assert(kv_size > 0);
uint32_t i = 0;
while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
++i;
}
assert(i < ubatch.n_tokens);
// Pad each stream to the reserve plan's block count.
const auto append_dummy_block = [&](llama_seq_id seq_id, uint32_t i) {
const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
const int32_t source_idx = state_source_idx(seq_id, ubatch.pos[i]);
const llama_pos pos = ubatch.pos[i];
const llama_seq_id seq_id = ubatch.seq_id[i][0];
const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
const int32_t source_idx = state_source_idx(seq_id, pos);
plan.state_write_idxs.push_back(cache_off + kv_size - 1);
plan.state_write_pos .push_back(0);
plan.state_write_idxs.push_back(cache_off + kv_size - 1);
plan.state_write_pos .push_back(0);
if (overlap) {
for (uint32_t j = 0; j < ratio; ++j) {
overlap_prev_reads.push_back(source_idx);
overlap_cur_reads .push_back(source_idx);
}
} else {
for (uint32_t j = 0; j < ratio; ++j) {
plan.state_read_idxs.push_back(source_idx);
}
}
};
if (overlap) {
for (uint32_t j = 0; j < ratio; ++j) {
overlap_prev_reads.push_back(source_idx);
overlap_cur_reads .push_back(source_idx);
if (dsv4_ubatch_has_coupled(ubatch)) {
if (plan.state_write_idxs.empty()) {
uint32_t i = 0;
while (i < ubatch.n_tokens && ubatch.pos[i] < 0) {
++i;
}
assert(i < ubatch.n_tokens);
append_dummy_block(ubatch.seq_id[i][0], i);
}
} else {
for (uint32_t j = 0; j < ratio; ++j) {
plan.state_read_idxs.push_back(source_idx);
const uint32_t n_blocks = (std::max<uint32_t>(1, ubatch.n_seq_tokens) + ratio - 1)/ratio;
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
const uint32_t n_writes = state_write_counts[seq_id];
if (n_writes >= n_blocks) {
continue;
}
if (n_writes + 1 != n_blocks) {
throw std::runtime_error("DSV4 CSA sequence positions are not contiguous");
}
uint32_t i = 0;
while (i < ubatch.n_tokens && (ubatch.pos[i] < 0 || !dsv4_token_has_seq(ubatch, i, seq_id))) {
++i;
}
assert(i < ubatch.n_tokens);
append_dummy_block(seq_id, i);
}
}
}
@@ -583,6 +620,63 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan(
plan.state_persist_dst_idxs.push_back(row.dst);
}
if (n_rs_seq > 0) {
for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {
const llama_seq_id seq_id = ubatch.seq_id_unq[s];
if (seq_id < 0 || (uint32_t) seq_id >= n_stream) {
continue;
}
const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);
const uint32_t rollback = (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0;
// Keep the restore graph fixed-width when no rollback is pending.
const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq ? (int64_t) rollback*state_rows : 0;
for (uint32_t r = 0; r < state_size; ++r) {
plan.state_restore_src_idxs.push_back((int32_t) (src_plane + stream_off + r));
plan.state_restore_dst_idxs.push_back((int32_t) (stream_off + r));
}
std::vector<uint32_t> token_idxs;
token_idxs.reserve(ubatch.n_tokens);
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
if (dsv4_token_has_seq(ubatch, i, seq_id)) {
token_idxs.push_back(i);
}
}
if (token_idxs.empty()) {
continue;
}
const uint32_t n_seq_tokens = (uint32_t) token_idxs.size();
const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq);
for (uint32_t d = 1; d <= n_rs_seq; ++d) {
const int64_t dst_plane = (int64_t) d*state_rows;
for (uint32_t r = 0; r < state_size; ++r) {
int32_t src;
if (d <= n_seq_tokens) {
const uint32_t prefix = n_seq_tokens - d;
src = (int32_t) (stream_off + r);
for (uint32_t j = 0; j < prefix; ++j) {
const uint32_t i_tok = token_idxs[j];
if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) {
src = (int32_t) (scratch_off + i_tok);
}
}
} else {
const int64_t src_plane = (int64_t) (d - n_seq_tokens)*state_rows;
src = (int32_t) (src_plane + stream_off + r);
}
plan.state_snapshot_src_idxs.push_back(src);
plan.state_snapshot_dst_idxs.push_back((int32_t) (dst_plane + stream_off + r));
}
}
}
}
static const bool debug = []() {
const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG");
return env && atoi(env) > 0;
@@ -604,12 +698,14 @@ static std::vector<llama_kv_cache_dsv4_context::comp_plan> dsv4_build_comp_plans
bool overlap,
uint32_t state_size,
uint32_t kv_size,
uint32_t n_stream) {
uint32_t n_stream,
uint32_t n_rs_seq,
const std::vector<uint32_t> & rs_idx) {
std::vector<llama_kv_cache_dsv4_context::comp_plan> plans;
plans.reserve(ubatches.size());
for (const llama_ubatch & ubatch : ubatches) {
plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream));
plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs_idx));
}
return plans;
@@ -696,7 +792,8 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(
bool overlap,
uint32_t state_size,
uint32_t kv_size,
uint32_t n_stream) {
uint32_t n_stream,
uint32_t n_rs_seq) {
llama_kv_cache_dsv4_context::comp_plan plan;
plan.n_visible.resize(ubatch.n_tokens);
plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream);
@@ -714,10 +811,16 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan(
const uint64_t state_rows = (uint64_t) state_size*n_stream;
const size_t n_persist = (size_t) std::min<uint64_t>(ubatch.n_tokens, state_rows);
const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq) : 0;
const size_t n_snapshot = (size_t) n_rs_seq*state_size*std::max<uint32_t>(1, ubatch.n_seqs_unq);
plan.state_pos .resize(ubatch.n_tokens);
plan.state_persist_src_idxs.resize(n_persist);
plan.state_persist_dst_idxs.resize(n_persist);
plan.state_restore_src_idxs.resize(n_restore);
plan.state_restore_dst_idxs.resize(n_restore);
plan.state_snapshot_src_idxs.resize(n_snapshot);
plan.state_snapshot_dst_idxs.resize(n_snapshot);
plan.state_read_idxs .resize((overlap ? 2u : 1u)*ratio*n_blocks);
plan.state_write_idxs.resize(n_blocks);
plan.state_write_pos .resize(n_blocks);
@@ -743,12 +846,14 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
uint32_t ratio,
uint32_t state_size,
uint32_t n_embd_state,
uint32_t n_rs_seq,
const char * name,
const llama_memory_i::layer_filter_cb & filter) :
ratio(ratio),
state_size(state_size),
n_embd_state(n_embd_state),
n_stream(unified ? 1 : n_seq_max) {
n_stream(unified ? 1 : n_seq_max),
n_rs_seq(n_rs_seq) {
const llama_hparams & hparams = model.hparams;
struct ggml_backend_buft_comparator {
@@ -804,8 +909,9 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
throw std::runtime_error("failed to create ggml context for DSV4 compressor state");
}
ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream);
ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream);
const uint32_t n_planes = n_stream*(1 + n_rs_seq);
ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);
ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes);
ggml_format_name(kv, "dsv4_%s_state_kv_l%d", name, il);
ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il);
@@ -837,8 +943,8 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
ctxs_bufs.emplace_back(std::move(ctx), buf);
}
LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, layers = %zu, size = %7.2f MiB\n",
__func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0);
LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, rs_seq = %u, layers = %zu, size = %7.2f MiB\n",
__func__, name, ratio, state_size, n_embd_state, n_stream, n_rs_seq, layers.size(), total_size()/1024.0/1024.0);
}
void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
@@ -848,9 +954,13 @@ void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
if (seq_id >= 0) {
GGML_ASSERT((uint32_t) seq_id < n_stream);
for (const auto & layer : layers) {
dsv4_clear_tensor_stream(layer.kv, (uint32_t) seq_id);
dsv4_clear_tensor_stream(layer.score, (uint32_t) seq_id);
for (uint32_t d = 0; d <= n_rs_seq; ++d) {
const uint32_t stream = d*n_stream + (uint32_t) seq_id;
dsv4_clear_tensor_stream(layer.kv, stream);
dsv4_clear_tensor_stream(layer.score, stream);
}
}
return;
}
@@ -868,6 +978,8 @@ void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_
return;
}
clear(seq_id_dst, true);
sc_info.ssrc.push_back((uint32_t) seq_id_src);
sc_info.sdst.push_back((uint32_t) seq_id_dst);
}
@@ -896,6 +1008,14 @@ uint32_t llama_dsv4_comp_state::get_n_stream() const {
return n_stream;
}
uint32_t llama_dsv4_comp_state::get_n_rs_seq() const {
return n_rs_seq;
}
uint32_t llama_dsv4_comp_state::get_n_rows() const {
return state_size*n_stream;
}
std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> ret;
for (const auto & [_, buf] : ctxs_bufs) {
@@ -905,13 +1025,26 @@ std::map<ggml_backend_buffer_type_t, size_t> llama_dsv4_comp_state::memory_break
return ret;
}
void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
void llama_dsv4_comp_state::state_write(
llama_io_write_i & io,
llama_seq_id seq_id,
llama_state_seq_flags flags,
const std::vector<uint32_t> & rs_idx) const {
GGML_UNUSED(flags);
uint32_t s0;
uint32_t ns;
dsv4_state_src_stream_range(n_stream, seq_id, s0, ns);
std::vector<uint32_t> stream_ids(ns);
for (uint32_t s = 0; s < ns; ++s) {
const uint32_t seq = seq_id >= 0 ? (uint32_t) seq_id : s0 + s;
if (seq >= rs_idx.size() || rs_idx[seq] > n_rs_seq) {
throw std::runtime_error("DSV4 recurrent state rollback index out of range");
}
stream_ids[s] = rs_idx[seq]*n_stream + s0 + s;
}
const uint32_t version = DSV4_COMP_STATE_VER;
const uint32_t n_layer = layers.size();
@@ -925,8 +1058,8 @@ void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_
for (const auto & layer : layers) {
io.write(&layer.il, sizeof(layer.il));
dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns);
dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns);
dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns, &stream_ids);
dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns, &stream_ids);
}
}
@@ -972,28 +1105,40 @@ void llama_dsv4_comp_state::state_read(llama_io_read_i & io, llama_seq_id seq_id
}
}
ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {
ggml_tensor * llama_dsv4_comp_state::get_kv_all(ggml_context * ctx, int32_t il) const {
const int32_t ids = map_layer_ids.at(il);
ggml_tensor * state = layers[ids].kv;
return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]);
return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);
}
ggml_tensor * llama_dsv4_comp_state::get_score_all(ggml_context * ctx, int32_t il) const {
const int32_t ids = map_layer_ids.at(il);
ggml_tensor * state = layers[ids].score;
return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0);
}
ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const {
ggml_tensor * state = get_kv_all(ctx, il);
const size_t row_size = ggml_row_size(state->type, state->ne[0]);
return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);
}
ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const {
const int32_t ids = map_layer_ids.at(il);
ggml_tensor * state = get_score_all(ctx, il);
const size_t row_size = ggml_row_size(state->type, state->ne[0]);
ggml_tensor * state = layers[ids].score;
return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]);
return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size);
}
ggml_tensor * llama_dsv4_comp_state::cpy_kv(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {
return ggml_set_rows(ctx, get_kv(ctx, il), cur, idxs);
return ggml_set_rows(ctx, get_kv_all(ctx, il), cur, idxs);
}
ggml_tensor * llama_dsv4_comp_state::cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const {
return ggml_set_rows(ctx, get_score(ctx, il), cur, idxs);
return ggml_set_rows(ctx, get_score_all(ctx, il), cur, idxs);
}
size_t llama_dsv4_comp_state::total_size() const {
@@ -1022,13 +1167,16 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
uint32_t n_seq_max,
uint32_t n_ubatch,
uint32_t n_pad,
uint32_t n_rs_seq,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse) :
hparams_raw(model.hparams),
hparams_csa(model.hparams),
hparams_hca(model.hparams),
hparams_lid(model.hparams),
n_seq_max(n_seq_max) {
n_seq_max(n_seq_max),
n_rs_seq(n_rs_seq),
rs_idx(n_seq_max, 0) {
const layer_filter_cb filter_raw = [&](int32_t il) {
if (filter && !filter(il)) {
@@ -1043,6 +1191,11 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
// Keep DSV4 KV/state streams per sequence even when public KV mode is unified.
const bool unified_raw = false;
hparams_raw.n_layer_nextn = 0;
hparams_csa.n_layer_nextn = 0;
hparams_hca.n_layer_nextn = 0;
hparams_lid.n_layer_nextn = 0;
LLAMA_LOG_INFO("%s: creating DSV4 raw KV cache\n", __func__);
dsv4_make_k_only(hparams_raw);
@@ -1109,19 +1262,19 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4(
csa_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
2*model.hparams.n_embd_head_k(), "csa", filter_csa);
2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa);
LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__);
hca_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO,
model.hparams.n_embd_head_k(), "hca", filter_hca);
model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca);
LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__);
lid_state = std::make_unique<llama_dsv4_comp_state>(
model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO,
2*model.hparams.indexer_head_size, "lid", filter_csa);
2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa);
// DSV4 attention reads compressed-K / compressor-state rows that the current
// graph does not necessarily overwrite; uninitialized buffer contents would
@@ -1255,17 +1408,35 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
}
if (p0 > 0) {
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max ||
p0 <= kv_raw->seq_pos_max(seq_id)) {
if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) {
return false;
}
bool res = true;
const llama_pos pos_max = kv_raw->seq_pos_max(seq_id);
if (p0 > pos_max) {
bool res = true;
res = res & kv_raw->seq_rm(seq_id, p0, -1);
res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
res = res & kv_raw->seq_rm(seq_id, p0, -1);
res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
return res;
}
if (n_rs_seq == 0) {
return false;
}
const llama_pos rollback = pos_max - (p0 - 1);
if (rollback < 1 || rollback > (llama_pos) n_rs_seq) {
return false;
}
const bool res = kv_raw->seq_rm(seq_id, p0, p1);
if (res) {
rs_idx[seq_id] = (uint32_t) rollback;
}
return res;
}
@@ -1290,6 +1461,10 @@ void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ds
csa_state->seq_cp(seq_id_src, seq_id_dst);
hca_state->seq_cp(seq_id_src, seq_id_dst);
lid_state->seq_cp(seq_id_src, seq_id_dst);
if (seq_id_src != seq_id_dst) {
rs_idx[seq_id_dst] = 0;
}
}
void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
@@ -1386,9 +1561,9 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id
dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid);
}
csa_state->state_write(io, seq_id, flags);
hca_state->state_write(io, seq_id, flags);
lid_state->state_write(io, seq_id, flags);
csa_state->state_write(io, seq_id, flags, rs_idx);
hca_state->state_write(io, seq_id, flags, rs_idx);
lid_state->state_write(io, seq_id, flags, rs_idx);
}
void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
@@ -1432,6 +1607,12 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id,
hca_state->state_read(io, seq_id, flags);
lid_state->state_read(io, seq_id, flags);
if (seq_id >= 0) {
GGML_ASSERT((uint32_t) seq_id < n_seq_max);
rs_idx[seq_id] = 0;
} else {
std::fill(rs_idx.begin(), rs_idx.end(), 0);
}
}
llama_kv_cache_iswa * llama_kv_cache_dsv4::get_raw() const {
@@ -1462,6 +1643,31 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const {
return lid_state.get();
}
uint32_t llama_kv_cache_dsv4::get_n_rs_seq() const {
return n_rs_seq;
}
const std::vector<uint32_t> & llama_kv_cache_dsv4::get_rs_idx() const {
return rs_idx;
}
void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches) {
if (n_rs_seq == 0) {
return;
}
for (const llama_ubatch & ubatch : ubatches) {
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) {
const llama_seq_id seq_id = ubatch.seq_id[i][s];
if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max) {
rs_idx[seq_id] = 0;
}
}
}
}
}
void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
if (seq_id < 0) {
kv_csa->clear(data);
@@ -1488,6 +1694,12 @@ void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) {
csa_state->clear(seq_id, data);
hca_state->clear(seq_id, data);
lid_state->clear(seq_id, data);
if (seq_id >= 0) {
rs_idx[seq_id] = 0;
} else {
std::fill(rs_idx.begin(), rs_idx.end(), 0);
}
}
//
@@ -1779,10 +1991,14 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
std::vector<llama_ubatch> ubatches_raw) :
ubatches(std::move(ubatches)),
plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream())),
kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream(),
kv->get_n_rs_seq(), kv->get_rs_idx())),
plans_hca(dsv4_build_comp_plans(this->ubatches, DSV4_HCA_RATIO, false,
kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream())),
plans_lid(plans_csa),
kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream(),
kv->get_n_rs_seq(), kv->get_rs_idx())),
plans_lid(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true,
kv->get_lid_state()->get_state_size(), kv->get_lid()->get_size(), kv->get_lid_state()->get_n_stream(),
kv->get_n_rs_seq(), kv->get_rs_idx())),
ctx_raw(std::make_unique<llama_kv_cache_dsv4_raw_context>(
kv->get_raw(),
std::move(sinfos_raw_base_write),
@@ -1809,6 +2025,7 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
hca_state(kv->get_hca_state()),
lid_state(kv->get_lid_state()),
status(ctx_raw->get_status()) {
kv->reset_rs_idx_for_ubatches(this->ubatches);
}
llama_kv_cache_dsv4_context::~llama_kv_cache_dsv4_context() = default;
@@ -1944,7 +2161,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
reserve_plan_csa = dsv4_build_reserve_comp_plan(
ubatch, DSV4_CSA_RATIO, true,
csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream());
csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream(), csa_state->get_n_rs_seq());
return reserve_plan_csa;
}
@@ -1958,7 +2175,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
reserve_plan_hca = dsv4_build_reserve_comp_plan(
ubatch, DSV4_HCA_RATIO, false,
hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream());
hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream(), hca_state->get_n_rs_seq());
return reserve_plan_hca;
}
@@ -1972,7 +2189,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_
reserve_plan_lid = dsv4_build_reserve_comp_plan(
ubatch, DSV4_CSA_RATIO, true,
lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream());
lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream(), lid_state->get_n_rs_seq());
return reserve_plan_lid;
}
src/llama-kv-cache-dsv4.h
+30 -5
View File
@@ -22,6 +22,7 @@ public:
uint32_t ratio,
uint32_t state_size,
uint32_t n_embd_state,
uint32_t n_rs_seq,
const char * name,
const llama_memory_i::layer_filter_cb & filter);
@@ -29,17 +30,21 @@ public:
void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst);
void apply_copies(const stream_copy_info & sc_info) const;
uint32_t get_ratio() const;
uint32_t get_ratio() const;
uint32_t get_state_size() const;
uint32_t get_n_stream() const;
uint32_t get_n_stream() const;
uint32_t get_n_rs_seq() const;
uint32_t get_n_rows() const;
std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const;
void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, const std::vector<uint32_t> & rs_idx) const;
void state_read (llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags);
ggml_tensor * get_kv (ggml_context * ctx, int32_t il) const;
ggml_tensor * get_score(ggml_context * ctx, int32_t il) const;
ggml_tensor * get_kv (ggml_context * ctx, int32_t il) const;
ggml_tensor * get_score (ggml_context * ctx, int32_t il) const;
ggml_tensor * get_kv_all (ggml_context * ctx, int32_t il) const;
ggml_tensor * get_score_all(ggml_context * ctx, int32_t il) const;
ggml_tensor * cpy_kv (ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const;
ggml_tensor * cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const;
@@ -59,6 +64,7 @@ private:
const uint32_t state_size;
const uint32_t n_embd_state;
const uint32_t n_stream;
const uint32_t n_rs_seq;
std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
@@ -93,6 +99,7 @@ public:
uint32_t n_seq_max,
uint32_t n_ubatch,
uint32_t n_pad,
uint32_t n_rs_seq,
const layer_filter_cb & filter,
const layer_reuse_cb & reuse);
@@ -141,6 +148,10 @@ public:
llama_dsv4_comp_state * get_hca_state() const;
llama_dsv4_comp_state * get_lid_state() const;
uint32_t get_n_rs_seq() const;
const std::vector<uint32_t> & get_rs_idx() const;
void reset_rs_idx_for_ubatches(const std::vector<llama_ubatch> & ubatches);
private:
llama_hparams hparams_raw;
llama_hparams hparams_csa;
@@ -148,6 +159,9 @@ private:
llama_hparams hparams_lid;
const uint32_t n_seq_max;
const uint32_t n_rs_seq;
std::vector<uint32_t> rs_idx;
std::unique_ptr<llama_kv_cache_iswa> kv_raw;
std::unique_ptr<llama_kv_cache> kv_csa;
@@ -268,6 +282,17 @@ public:
std::vector<int32_t> state_persist_src_idxs;
std::vector<int32_t> state_persist_dst_idxs;
// Device-side rollback restore copies snapshot planes back to the
// current compressor-state plane before the graph reads it.
std::vector<int32_t> state_restore_src_idxs;
std::vector<int32_t> state_restore_dst_idxs;
// Device-side rollback snapshots copy rows from the graph-local
// [persistent_state | current_ubatch_scratch] tensor into rollback
// planes after the graph has computed current-token compressor state.
std::vector<int32_t> state_snapshot_src_idxs;
std::vector<int32_t> state_snapshot_dst_idxs;
// Flattened source row ids used for state-backed commits. Source rows
// index the graph-local [persistent_state | current_ubatch_scratch]
// tensor. For overlapped compression the first half is previous rows
src/llama-model-loader.cpp
+2
View File
@@ -526,6 +526,7 @@ llama_model_loader::llama_model_loader(
llama_load_mode load_mode,
bool check_tensors,
bool no_alloc,
bool load_mtp,
const llama_model_kv_override * param_overrides_p,
const llama_model_tensor_buft_override * param_tensor_buft_overrides_p)
: metadata(meta), set_tensor_data(set_tensor_data), set_tensor_data_ud(set_tensor_data_ud) {
@@ -813,6 +814,7 @@ llama_model_loader::llama_model_loader(
this->check_tensors = check_tensors;
this->no_alloc = no_alloc;
this->load_mtp = load_mtp;
}
std::string llama_model_loader::get_arch_name() const {
src/llama-model-loader.h
+2
View File
@@ -79,6 +79,7 @@ struct llama_model_loader {
bool use_direct_io = false;
bool check_tensors;
bool no_alloc;
bool load_mtp;
llama_files files;
llama_ftype ftype;
@@ -129,6 +130,7 @@ struct llama_model_loader {
llama_load_mode load_mode,
bool check_tensors,
bool no_alloc,
bool load_mtp,
const llama_model_kv_override * param_overrides_p,
const llama_model_tensor_buft_override * param_tensor_buft_overrides_p);
src/llama-model.cpp
+83 -43
View File
@@ -2213,24 +2213,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
{
res = nullptr;
} break;
case LLM_ARCH_DEEPSEEK32:
{
res = new llama_kv_cache_dsa(
*this,
params.type_k,
params.type_v,
!cparams.flash_attn,
cparams.offload_kqv,
cparams.kv_unified,
cparams.n_ctx_seq,
cparams.n_seq_max,
1,
hparams.n_swa,
hparams.swa_type,
nullptr,
nullptr);
} break;
case LLM_ARCH_GLM_DSA:
case LLM_ARCH_DEEPSEEK32:
{
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) {
// The NextN/MTP draft head runs dense MLA (no DSA indexer), so the
@@ -2280,15 +2264,84 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
nullptr);
}
} break;
case LLM_ARCH_DEEPSEEK4:
{
GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
const llama_memory_i::layer_filter_cb filter_mtp = [&](int32_t il) {
return il >= (int32_t) hparams.n_layer();
};
res = new llama_kv_cache_iswa(
*this,
params.type_k,
params.type_v,
!cparams.flash_attn,
cparams.offload_kqv,
params.swa_full,
cparams.kv_unified,
cparams.n_ctx_seq,
cparams.n_seq_max,
cparams.n_ubatch,
1,
nullptr,
filter_mtp,
nullptr,
nullptr);
} else {
res = new llama_kv_cache_dsv4(
*this,
params.type_k,
params.type_v,
!cparams.flash_attn,
cparams.offload_kqv,
params.swa_full,
cparams.kv_unified,
cparams.n_ctx_seq,
cparams.n_seq_max,
cparams.n_ubatch,
1,
cparams.n_rs_seq,
nullptr,
nullptr);
}
} break;
case LLM_ARCH_DFLASH:
{
// DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring)
if (hparams.dsv4_hc_mult > 0) {
GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
res = new llama_kv_cache_iswa(
*this,
params.type_k,
params.type_v,
!cparams.flash_attn,
cparams.offload_kqv,
params.swa_full,
cparams.kv_unified,
cparams.n_ctx_seq,
cparams.n_seq_max,
cparams.n_ubatch,
1,
nullptr,
nullptr,
nullptr,
nullptr);
break;
}
}
[[fallthrough]];
// Models that need standard caching should rely on recurrent/hybrid
// checks
default:
{
// The MTP head is dense-attention only on hybrid Qwen3.5/3.6, so use a plain
// The MTP head is dense-attention only on hybrid Qwen3-Next/3.5/3.6, so use a plain
// attention KV cache for the MTP context instead of the hybrid wrapper.
const bool mtp_on_hybrid_qwen35 =
const bool mtp_on_hybrid_qwen =
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
(arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
(arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
if (llm_arch_is_recurrent(arch)) {
res = new llama_memory_recurrent(
@@ -2300,7 +2353,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
cparams.n_seq_max,
cparams.n_rs_seq,
nullptr);
} else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen35) {
} else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) {
// The main difference between hybrid architectures is the
// layer filters, so pick the right one here
llama_memory_hybrid::layer_filter_cb filter_attn = nullptr;
@@ -2315,7 +2368,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
filter_recr = [&](uint32_t il) {
return hparams.is_recr(il) && hparams.n_ff(il) == 0;
};
} else if (arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
} else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) {
filter_attn = [&](uint32_t il) {
return il < hparams.n_layer() && !hparams.is_recr(il);
};
@@ -2381,12 +2434,12 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
};
}
if (mtp_on_hybrid_qwen35) {
if (mtp_on_hybrid_qwen) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
}
if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA ||
arch == LLM_ARCH_MIMO2) &&
arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_DEEPSEEK32) &&
hparams.n_layer_nextn > 0) {
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
@@ -2395,24 +2448,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
}
}
if (arch == LLM_ARCH_DEEPSEEK4) {
GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
res = new llama_kv_cache_dsv4(
*this,
params.type_k,
params.type_v,
!cparams.flash_attn,
cparams.offload_kqv,
params.swa_full,
cparams.kv_unified,
cparams.n_ctx_seq,
cparams.n_seq_max,
cparams.n_ubatch,
1,
filter,
reuse);
} else if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
GGML_ASSERT(hparams.is_swa_any());
if (arch == LLM_ARCH_GEMMA4_ASSISTANT) {
@@ -2532,6 +2568,7 @@ llama_model_params llama_model_default_params() {
/*.use_extra_bufts =*/ true,
/*.no_host =*/ false,
/*.no_alloc =*/ false,
/*.load_mtp =*/ false,
};
return result;
@@ -2760,9 +2797,12 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_STEP35:
case LLM_ARCH_TALKIE:
case LLM_ARCH_MELLUM:
case LLM_ARCH_DFLASH:
return LLAMA_ROPE_TYPE_NEOX;
case LLM_ARCH_DFLASH:
// DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX
return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX;
case LLM_ARCH_QWEN2VL:
case LLM_ARCH_PADDLEOCR:
return LLAMA_ROPE_TYPE_MROPE;
src/llama-quant.cpp
+1 -1
View File
@@ -895,7 +895,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
const llama_model_kv_override * kv_overrides = params->kv_overrides;
std::vector<std::string> splits = {};
llama_model_loader ml(/*metadata*/ nullptr, /*set_tensor_data*/ nullptr, /*set_tensor_data_ud*/ nullptr,
fname_inp, splits, /*file*/ nullptr, /*load_mode*/ load_mode, /*check_tensors*/ true, /*no_alloc*/ false, kv_overrides, nullptr);
fname_inp, splits, /*file*/ nullptr, /*load_mode*/ load_mode, /*check_tensors*/ true, /*no_alloc*/ false, /*load_mtp*/ true, kv_overrides, nullptr);
ml.init_mappings(false); // no prefetching
auto mparams = llama_model_default_params();
src/llama.cpp
+1 -1
View File
@@ -325,7 +325,7 @@ static std::pair<int, llama_model *> llama_model_load(struct gguf_context * meta
const std::string & fname, std::vector<std::string> & splits, FILE * file, llama_model_params & params) {
try {
llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode,
params.check_tensors, params.no_alloc, params.kv_overrides, params.tensor_buft_overrides);
params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides);
ml.print_info();
std::unique_ptr<llama_model> model_ptr(llama_model_create(ml, params));
src/models/cohere2moe.cpp
+5 -1
View File
@@ -55,7 +55,11 @@ void llama_model_cohere2moe::load_arch_tensors(llama_model_loader & ml) {
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
src/models/deepseek32.cpp
+269 -10
View File
@@ -44,13 +44,24 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) {
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer");
switch (hparams.n_layer()) {
case 62: type = LLM_TYPE_685B_A37B; break;
case 61: type = LLM_TYPE_685B_A37B; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
void llama_model_deepseek32::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const bool is_mla = hparams.is_mla();
if (!is_mla) {
throw std::runtime_error("DEEPSEEK32 architecture requires MLA");
@@ -80,12 +91,7 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
}
for (int i = 0; i < n_layer_all; ++i) {
int flags = 0;
if (i >= n_layer) {
// skip all tensors in the NextN layers
// TODO @ngxson : TENSOR_NOT_REQUIRED was a hack, need to remove it later
flags |= TENSOR_SKIP | TENSOR_NOT_REQUIRED;
}
const int flags = (i >= n_layer) ? mtp_flags : trunk_flags;
auto & layer = layers[i];
@@ -138,7 +144,7 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
}
// NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers
// NextN/MTP tensors - conditionally load for last nextn_predict_layers
if (i >= n_layer) {
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags);
@@ -153,6 +159,9 @@ void llama_model_deepseek32::load_arch_tensors(llama_model_loader &) {
}
std::unique_ptr<llm_graph_context> llama_model_deepseek32::build_arch_graph(const llm_graph_params & params) const {
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
return std::make_unique<graph_mtp>(*this, params);
}
return std::make_unique<graph>(*this, params);
}
@@ -430,7 +439,9 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il);
}
}
if (il == n_layer - 1 && inp_out_ids) {
// when unmasked nextn embeddings are requested, t_h_nextn must keep all rows,
// so the early output masking has to be skipped (it is applied after the final norm instead)
if (il == n_layer - 1 && inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -493,6 +504,14 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
// post-norm hidden state feeds the NextN/MTP draft head
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cb(cur, "result_norm", -1);
res->t_embd = cur;
@@ -504,3 +523,243 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
ggml_build_forward_expand(gf, cur);
}
// LLM_GRAPH_TYPE_DECODER_MTP draft head for DeepSeek V3.2 (DEEPSEEK32).
// Semantics mirror the deepseek-family NextN/MTP layer:
// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj ->
// full deepseek32 decoder block (dense MLA attention + sigmoid-gated MoE FFN
// with shared expert, exactly as the trunk deepseek2 graph builds it) ->
// shared_head_norm (fallback output_norm) -> shared LM head.
// The DSA indexer is not used at runtime.
llama_model_deepseek32::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
GGML_ASSERT(hparams.n_layer_nextn > 0 && "DEEPSEEK32 MTP requires n_layer_nextn > 0");
GGML_ASSERT(hparams.n_layer_nextn == 1 && "DEEPSEEK32 MTP currently only supports a single MTP block");
GGML_ASSERT(hparams.is_mla() && "DEEPSEEK32 MTP requires MLA");
const int il = hparams.n_layer() + cparams.nextn_layer_offset;
GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
"nextn_layer_offset out of range [0, n_layer_nextn)");
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");
// note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA
const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();
const int64_t n_embd_head_qk_rope = hparams.n_rot();
const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;
const uint32_t kv_lora_rank = hparams.n_lora_kv;
// We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.
// See the deepseek2 trunk graph for the detailed explanation - this must match it EXACTLY.
GGML_ASSERT(ext_factor >= 0.0f);
const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));
const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));
const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));
// TODO: extract in a common llm_graph_context::build_inp_embd_h()
auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
ggml_tensor * tok_embd;
if (ubatch.token) {
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
} else {
tok_embd = inp->embd;
}
cb(tok_embd, "mtp_tok_embd", il);
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
ggml_set_input(inp->h);
ggml_set_name(inp->h, "mtp_h_input");
ggml_tensor * h_embd = inp->h;
res->add_input(std::move(inp));
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
// MLA with the absorption optimization uses a K-only cache (V is a view of K)
auto * inp_attn = build_attn_inp_k();
ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
cb(e_norm, "mtp_enorm", il);
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
cb(concat, "mtp_concat", il);
ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
cb(cur, "mtp_eh_proj", il);
ggml_tensor * inpSA = cur;
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_norm", il);
// self-attention: dense MLA, same construction as the deepseek2 trunk graph
{
ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur);
cb(q, "mtp_q", il);
q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
cb(q, "mtp_q", il);
q = ggml_mul_mat(ctx0, layer.wq_b, q);
cb(q, "mtp_q", il);
// split into {n_embd_head_qk_nope, n_head, n_tokens}
ggml_tensor * q_nope =
ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),
ggml_row_size(q->type, n_embd_head_k) * n_head, 0);
cb(q_nope, "mtp_q_nope", il);
// and {n_embd_head_qk_rope, n_head, n_tokens}
ggml_tensor * q_pe = ggml_view_3d(
ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),
ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));
cb(q_pe, "mtp_q_pe", il);
ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
cb(kv_cmpr_pe, "mtp_kv_cmpr_pe", il);
// split into {kv_lora_rank, n_tokens}
ggml_tensor * kv_cmpr =
ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);
cb(kv_cmpr, "mtp_kv_cmpr", il);
// and {n_embd_head_qk_rope, 1, n_tokens}
ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),
ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
cb(k_pe, "mtp_k_pe", il);
q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(q_pe, "mtp_q_pe", il);
k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(k_pe, "mtp_k_pe", il);
kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
cb(kv_cmpr, "mtp_kv_cmpr", il);
// {n_embd_head_qk_nope, n_tokens, n_head}
q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
cb(q_nope, "mtp_q_nope_perm", il);
// {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}
ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
cb(q_nope_absorbed, "mtp_q_nope_absorbed", il);
// {kv_lora_rank, n_head, n_tokens}
q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
cb(q_nope_absorbed, "mtp_q_nope_absorbed_perm", il);
// {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}
// note: rope must go first for in-place context shifting in build_rope_shift()
ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
cb(Qcur, "mtp_Qcur", il);
kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
cb(kv_cmpr, "mtp_kv_cmpr_reshape", il);
// {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}
ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);
cb(Kcur, "mtp_Kcur", il);
// {kv_lora_rank, 1, n_tokens}
ggml_tensor * Vcur = kv_cmpr;
cb(Vcur, "mtp_Vcur", il);
// note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)
cur = build_attn(inp_attn,
layer.wo, NULL, layer.wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale, il);
cb(cur, "mtp_attn_out", il);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "mtp_ffn_inp", il);
cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "mtp_ffn_norm", il);
// MoE FFN with shared expert - same construction as the deepseek2 trunk graph
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, n_expert_used,
LLM_FFN_SILU, hparams.expert_weights_norm,
hparams.expert_weights_scale,
(llama_expert_gating_func_type) hparams.expert_gating_func,
il,
nullptr,
layer.ffn_gate_up_exps,
layer.ffn_up_exps_s,
layer.ffn_gate_exps_s,
layer.ffn_down_exps_s);
cb(moe_out, "mtp_ffn_moe_out", il);
// FFN shared expert
ggml_tensor * ffn_shexp =
build_ffn(cur,
layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
layer.ffn_gate_shexp, NULL, layer.ffn_gate_shexp_s,
layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(ffn_shexp, "mtp_ffn_shexp", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
cb(cur, "mtp_ffn_out", il);
cur = ggml_add(ctx0, cur, ffn_inp);
cb(cur, "mtp_post_ffn", il);
// shared_head_norm applied after the decoder block, before the shared LM head.
// The post-norm hidden state seeds the next MTP step.
ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
? layer.nextn.shared_head_norm
: model.output_norm;
GGML_ASSERT(head_norm_w && "DEEPSEEK32 MTP: missing both nextn.shared_head_norm and output_norm");
cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
cb(cur, "mtp_shared_head_norm", -1);
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
GGML_ASSERT(head_w && "DEEPSEEK32 MTP: missing LM head (nextn.shared_head_head or model.output)");
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
src/models/deepseek4.cpp
+411 -70
View File
@@ -16,6 +16,16 @@ static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) {
}
void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
if (hparams.n_layer_nextn > 0 && hparams.n_layer_nextn < hparams.n_layer_all) {
const uint32_t n_layer_main = hparams.n_layer_all - hparams.n_layer_nextn;
const std::string mtp_probe = "blk." + std::to_string(n_layer_main) + ".nextn.eh_proj.weight";
if (ml.get_weight(mtp_probe.c_str()) == nullptr) {
hparams.n_layer_nextn = 0;
}
}
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < block_count");
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
@@ -24,8 +34,8 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
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_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer());
if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer(), 0)) {
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;
}
@@ -41,9 +51,11 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
ml.get_key(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count);
hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd;
uint32_t n_compress_ratios = 0;
ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_compress_ratios);
if (n_compress_ratios < hparams.n_layer()) {
if (n_compress_ratios < hparams.n_layer_all) {
throw std::runtime_error("DeepSeek-V4 compress_ratios is shorter than block_count");
}
ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios);
@@ -54,6 +66,9 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
}
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
hparams.set_swa_pattern(0);
for (uint32_t il = hparams.n_layer(); il < hparams.n_layer_all; ++il) {
hparams.is_swa_impl[il] = true;
}
switch (hparams.n_layer()) {
case 43: type = LLM_TYPE_UNKNOWN; break;
@@ -61,7 +76,7 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) {
}
}
void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) {
void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const int64_t q_lora_rank = hparams.n_lora_q;
@@ -75,6 +90,10 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) {
const int64_t hc_dim = hc_mult * n_embd;
const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = ml.load_mtp ? 0 : TENSOR_SKIP;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
@@ -84,69 +103,82 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) {
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) {
for (int i = 0; i < n_layer_all; ++i) {
auto & layer = layers[i];
const int flags = i < n_layer ? trunk_flags : mtp_flags;
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.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, flags);
layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags);
layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags);
layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, flags);
layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, flags);
layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, flags);
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}, flags);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, flags);
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.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, flags);
layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, flags);
layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, flags);
layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, flags);
layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, flags);
layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, flags);
const int64_t ratio = hparams.dsv4_compress_ratios[i];
if (ratio != 0) {
const int64_t coff = ratio == 4 ? 2 : 1;
layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, coff * n_embd_head}, 0);
layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, 0);
layer.attn_comp_ape = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE, "weight", i), {coff * n_embd_head, ratio}, 0);
layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, 0);
layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, coff * n_embd_head}, flags);
layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, flags);
layer.attn_comp_ape = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE, "weight", i), {coff * n_embd_head, ratio}, flags);
layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, flags);
if (ratio == 4) {
const int64_t n_embd_indexer = hparams.indexer_head_size;
layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, 0);
layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, 0);
layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags);
layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, flags);
layer.indexer_comp_wkv = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "weight", i), {n_embd, 2 * n_embd_indexer}, 0);
layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, 0);
layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", i), {2 * n_embd_indexer, ratio}, 0);
layer.indexer_comp_norm = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "weight", i), {n_embd_indexer}, 0);
layer.indexer_comp_wkv = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "weight", i), {n_embd, 2 * n_embd_indexer}, flags);
layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, flags);
layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", i), {2 * n_embd_indexer, ratio}, flags);
layer.indexer_comp_norm = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "weight", i), {n_embd_indexer}, flags);
} else if (ratio != 128) {
throw std::runtime_error("DeepSeek-V4 loader only supports compression ratios 0, 4, and 128");
}
}
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags);
if ((uint32_t) i < hparams.dsv4_hash_layer_count) {
layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, 0);
layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, flags);
} else {
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);
layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags);
}
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
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_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags);
layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags);
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);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
if (i >= n_layer) {
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags);
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags);
layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED | flags);
}
}
}
std::unique_ptr<llm_graph_context> llama_model_deepseek4::build_arch_graph(const llm_graph_params & params) const {
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
return std::make_unique<graph_mtp>(*this, params);
}
return std::make_unique<graph>(*this, params);
}
@@ -175,18 +207,69 @@ static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, b
return ggml_concat(ctx, t, row, 1);
}
static ggml_tensor * dsv4_with_zero_dep(ggml_context * ctx, ggml_tensor * t, ggml_tensor * dep) {
if (dep == nullptr) {
return t;
struct dsv4_state_tensors {
ggml_tensor * kv;
ggml_tensor * score;
};
static dsv4_state_tensors dsv4_build_state_restore(
ggml_context * ctx,
const llm_graph_input_dsv4::comp_input & inp,
const llama_dsv4_comp_state * state,
int32_t il) {
dsv4_state_tensors restored = {
state->get_kv_all(ctx, il),
state->get_score_all(ctx, il),
};
if (inp.state_restore_src_idxs == nullptr || inp.state_restore_dst_idxs == nullptr) {
return restored;
}
ggml_tensor * zero = ggml_scale(ctx, ggml_sum(ctx, dep), 0.0f);
return ggml_add(ctx, t, zero);
ggml_tensor * kv_rows = ggml_get_rows(ctx, restored.kv, inp.state_restore_src_idxs);
restored.kv = state->cpy_kv(ctx, kv_rows, inp.state_restore_dst_idxs, il);
ggml_tensor * score_rows = ggml_get_rows(ctx, restored.score, inp.state_restore_src_idxs);
restored.score = state->cpy_score(ctx, score_rows, inp.state_restore_dst_idxs, il);
return restored;
}
static dsv4_state_tensors dsv4_build_state_snapshot(
ggml_context * ctx,
const llm_graph_input_dsv4::comp_input & inp,
const llama_dsv4_comp_state * state,
ggml_tensor * source_kv,
ggml_tensor * source_score,
int32_t il) {
if (inp.state_snapshot_src_idxs == nullptr || inp.state_snapshot_dst_idxs == nullptr ||
source_kv == nullptr || source_score == nullptr) {
return {};
}
ggml_tensor * kv_rows = ggml_get_rows(ctx, source_kv, inp.state_snapshot_src_idxs);
ggml_tensor * kv = state->cpy_kv(ctx, kv_rows, inp.state_snapshot_dst_idxs, il);
ggml_tensor * score_rows = ggml_get_rows(ctx, source_score, inp.state_snapshot_src_idxs);
ggml_tensor * score = state->cpy_score(ctx, score_rows, inp.state_snapshot_dst_idxs, il);
return { kv, score };
}
static constexpr int64_t DSV4_CSA_RATIO = 4;
static constexpr int64_t DSV4_HCA_RATIO = 128;
// mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens]
static ggml_tensor * dsv4_hc_mean(ggml_context * ctx, ggml_tensor * x) {
const int64_t hc = x->ne[1];
ggml_tensor * acc = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0);
for (int64_t s = 1; s < hc; ++s) {
acc = ggml_add(ctx, acc, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1]));
}
return ggml_scale(ctx, acc, 1.0f/hc);
}
static ggml_tensor * dsv4_hc_affine(
ggml_context * ctx,
ggml_tensor * x,
@@ -804,8 +887,29 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
ggml_tensor * cur,
ggml_tensor * inp_pos,
int il) const {
return build_attention_impl(model, inp_dsv4, nullptr, cur, inp_pos, il);
}
ggml_tensor * llama_model_deepseek4::graph::build_attention(
const llama_model & model,
llm_graph_input_attn_k_iswa * inp_mtp,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int il) const {
return build_attention_impl(model, nullptr, inp_mtp, cur, inp_pos, il);
}
ggml_tensor * llama_model_deepseek4::graph::build_attention_impl(
const llama_model & model,
llm_graph_input_dsv4 * inp_dsv4,
llm_graph_input_attn_k_iswa * inp_mtp,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int il) const {
GGML_ASSERT((inp_dsv4 == nullptr) != (inp_mtp == nullptr));
const auto & layer = model.layers[il];
llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw();
llm_graph_input_dsv4_raw * inp_attn = inp_dsv4 ? inp_dsv4->get_raw() : nullptr;
const int64_t n_embd_head = hparams.n_embd_head_k();
const int64_t n_embd_head_rope = hparams.n_rot();
@@ -873,9 +977,12 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
cb(kv, "kv", il);
const int64_t ratio = hparams.dsv4_compress_ratios[il];
GGML_ASSERT(inp_dsv4 || ratio == 0);
ggml_tensor * hca_state_kv = nullptr;
ggml_tensor * hca_state_score = nullptr;
ggml_tensor * hca_source_kv = nullptr;
ggml_tensor * hca_source_score = nullptr;
if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) {
hca_state_kv = build_lora_mm(layer.attn_comp_wkv, cur);
cb(hca_state_kv, "hca_state_kv", il);
@@ -906,10 +1013,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
GGML_ASSERT(inp_dsv4->get_csa().state_write_idxs);
ggml_tensor * csa_source_kv = ggml_concat(ctx0,
inp_dsv4->mctx->get_csa_state()->get_kv(ctx0, il), csa_state_kv, 1);
ggml_tensor * csa_source_score = ggml_concat(ctx0,
inp_dsv4->mctx->get_csa_state()->get_score(ctx0, il), csa_state_score, 1);
const auto * csa_state = inp_dsv4->mctx->get_csa_state();
const dsv4_state_tensors csa_restored = dsv4_build_state_restore(
ctx0, inp_dsv4->get_csa(), csa_state, il);
ggml_tensor * csa_base_kv = dsv4_view_2d(
ctx0, csa_restored.kv, csa_restored.kv->ne[0], csa_state->get_n_rows(), 0);
ggml_tensor * csa_base_score = dsv4_view_2d(
ctx0, csa_restored.score, csa_restored.score->ne[0], csa_state->get_n_rows(), 0);
ggml_tensor * csa_source_kv = ggml_concat(ctx0, csa_base_kv, csa_state_kv, 1);
ggml_tensor * csa_source_score = ggml_concat(ctx0, csa_base_score, csa_state_score, 1);
ggml_tensor * kv_comp_csa_state = build_overlap_compressed_kv_from_state(
csa_source_kv,
@@ -930,8 +1043,19 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0,
kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il));
csa_state_kv = dsv4_with_zero_dep(ctx0, csa_state_kv, kv_comp_csa_state);
csa_state_score = dsv4_with_zero_dep(ctx0, csa_state_score, kv_comp_csa_state);
ggml_tensor * csa_snapshot_source_kv = ggml_concat(ctx0,
csa_restored.kv, csa_state_kv, 1);
ggml_tensor * csa_snapshot_source_score = ggml_concat(ctx0,
csa_restored.score, csa_state_score, 1);
const dsv4_state_tensors csa_snapshot = dsv4_build_state_snapshot(
ctx0, inp_dsv4->get_csa(), csa_state, csa_snapshot_source_kv, csa_snapshot_source_score, il);
if (csa_snapshot.kv != nullptr) {
ggml_build_forward_expand(gf, csa_snapshot.kv);
}
if (csa_snapshot.score != nullptr) {
ggml_build_forward_expand(gf, csa_snapshot.score);
}
ggml_tensor * csa_persist_kv = ggml_get_rows(ctx0, csa_state_kv, inp_dsv4->get_csa().state_persist_src_idxs);
ggml_tensor * csa_persist_score = ggml_get_rows(ctx0, csa_state_score, inp_dsv4->get_csa().state_persist_src_idxs);
@@ -958,10 +1082,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
GGML_ASSERT(inp_dsv4->get_lid().state_write_idxs);
ggml_tensor * lid_source_kv = ggml_concat(ctx0,
inp_dsv4->mctx->get_lid_state()->get_kv(ctx0, il), lid_state_kv, 1);
ggml_tensor * lid_source_score = ggml_concat(ctx0,
inp_dsv4->mctx->get_lid_state()->get_score(ctx0, il), lid_state_score, 1);
const auto * lid_state = inp_dsv4->mctx->get_lid_state();
const dsv4_state_tensors lid_restored = dsv4_build_state_restore(
ctx0, inp_dsv4->get_lid(), lid_state, il);
ggml_tensor * lid_base_kv = dsv4_view_2d(
ctx0, lid_restored.kv, lid_restored.kv->ne[0], lid_state->get_n_rows(), 0);
ggml_tensor * lid_base_score = dsv4_view_2d(
ctx0, lid_restored.score, lid_restored.score->ne[0], lid_state->get_n_rows(), 0);
ggml_tensor * lid_source_kv = ggml_concat(ctx0, lid_base_kv, lid_state_kv, 1);
ggml_tensor * lid_source_score = ggml_concat(ctx0, lid_base_score, lid_state_score, 1);
ggml_tensor * kv_comp_lid_state = build_overlap_compressed_kv_from_state(
lid_source_kv,
@@ -982,8 +1112,19 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
ggml_build_forward_expand(gf, inp_dsv4->mctx->get_lid()->cpy_k(ctx0,
kv_comp_lid_state, inp_dsv4->get_lid().state_write_idxs, il));
lid_state_kv = dsv4_with_zero_dep(ctx0, lid_state_kv, kv_comp_lid_state);
lid_state_score = dsv4_with_zero_dep(ctx0, lid_state_score, kv_comp_lid_state);
ggml_tensor * lid_snapshot_source_kv = ggml_concat(ctx0,
lid_restored.kv, lid_state_kv, 1);
ggml_tensor * lid_snapshot_source_score = ggml_concat(ctx0,
lid_restored.score, lid_state_score, 1);
const dsv4_state_tensors lid_snapshot = dsv4_build_state_snapshot(
ctx0, inp_dsv4->get_lid(), lid_state, lid_snapshot_source_kv, lid_snapshot_source_score, il);
if (lid_snapshot.kv != nullptr) {
ggml_build_forward_expand(gf, lid_snapshot.kv);
}
if (lid_snapshot.score != nullptr) {
ggml_build_forward_expand(gf, lid_snapshot.score);
}
ggml_tensor * lid_persist_kv = ggml_get_rows(ctx0, lid_state_kv, inp_dsv4->get_lid().state_persist_src_idxs);
ggml_tensor * lid_persist_score = ggml_get_rows(ctx0, lid_state_score, inp_dsv4->get_lid().state_persist_src_idxs);
@@ -997,15 +1138,21 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
ggml_build_forward_expand(gf, lid_state_score);
}
ggml_tensor * hca_state_dep = nullptr;
const llama_dsv4_comp_state * hca_state = nullptr;
dsv4_state_tensors hca_restored = {};
if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_write_idxs) {
GGML_ASSERT(hca_state_kv);
GGML_ASSERT(hca_state_score);
ggml_tensor * hca_source_kv = ggml_concat(ctx0,
inp_dsv4->mctx->get_hca_state()->get_kv(ctx0, il), hca_state_kv, 1);
ggml_tensor * hca_source_score = ggml_concat(ctx0,
inp_dsv4->mctx->get_hca_state()->get_score(ctx0, il), hca_state_score, 1);
hca_state = inp_dsv4->mctx->get_hca_state();
hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il);
ggml_tensor * hca_base_kv = dsv4_view_2d(
ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0);
ggml_tensor * hca_base_score = dsv4_view_2d(
ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0);
hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1);
hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1);
ggml_tensor * kv_comp_hca = build_hca_compressed_kv_from_state(
hca_source_kv,
@@ -1024,15 +1171,41 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0,
kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il));
hca_state_dep = kv_comp_hca;
}
if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) {
GGML_ASSERT(hca_state_kv);
GGML_ASSERT(hca_state_score);
hca_state_kv = dsv4_with_zero_dep(ctx0, hca_state_kv, hca_state_dep);
hca_state_score = dsv4_with_zero_dep(ctx0, hca_state_score, hca_state_dep);
if (hca_state == nullptr) {
hca_state = inp_dsv4->mctx->get_hca_state();
}
if (hca_restored.kv == nullptr) {
hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il);
}
if (hca_source_kv == nullptr || hca_source_score == nullptr) {
ggml_tensor * hca_base_kv = dsv4_view_2d(
ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0);
ggml_tensor * hca_base_score = dsv4_view_2d(
ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0);
hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1);
hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1);
}
ggml_tensor * hca_snapshot_source_kv = ggml_concat(ctx0,
hca_restored.kv, hca_state_kv, 1);
ggml_tensor * hca_snapshot_source_score = ggml_concat(ctx0,
hca_restored.score, hca_state_score, 1);
const dsv4_state_tensors hca_snapshot = dsv4_build_state_snapshot(
ctx0, inp_dsv4->get_hca(), hca_state, hca_snapshot_source_kv, hca_snapshot_source_score, il);
if (hca_snapshot.kv != nullptr) {
ggml_build_forward_expand(gf, hca_snapshot.kv);
}
if (hca_snapshot.score != nullptr) {
ggml_build_forward_expand(gf, hca_snapshot.score);
}
ggml_tensor * hca_persist_kv = ggml_get_rows(ctx0, hca_state_kv, inp_dsv4->get_hca().state_persist_src_idxs);
ggml_tensor * hca_persist_score = ggml_get_rows(ctx0, hca_state_score, inp_dsv4->get_hca().state_persist_src_idxs);
@@ -1047,7 +1220,14 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention(
}
ggml_tensor * out = nullptr;
if (ratio == DSV4_CSA_RATIO &&
if (inp_mtp) {
out = build_attn(inp_mtp,
nullptr, nullptr, nullptr,
q, kv, nullptr,
nullptr, layer.attn_sinks, nullptr,
1.0f/sqrtf(float(n_embd_head)), il);
cb(out, "attn_raw", il);
} else if (ratio == DSV4_CSA_RATIO &&
inp_dsv4->get_csa().kq_mask &&
inp_dsv4->get_lid().kq_mask &&
inp_dsv4->get_lid().k_rot) {
@@ -1106,6 +1286,12 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
cb(inpL, "hc_init", -1);
for (int il = 0; il < n_layer; ++il) {
if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) {
res->t_layer_inp[il] = dsv4_hc_mean(ctx0, inpL);
cb(res->t_layer_inp[il], "layer_inp", il);
ggml_build_forward_expand(gf, res->t_layer_inp[il]);
}
ggml_tensor * residual = inpL;
ggml_tensor * post = nullptr;
ggml_tensor * comb = nullptr;
@@ -1182,10 +1368,23 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
cb(inpL, "l_last", il);
}
if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) {
res->t_layer_inp[n_layer] = dsv4_hc_mean(ctx0, inpL);
cb(res->t_layer_inp[n_layer], "layer_inp", n_layer);
ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]);
}
ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
ggml_tensor * flat_out = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat;
if (cparams.embeddings_nextn) {
ggml_tensor * h_nextn = cparams.embeddings_nextn_masked ? flat_out : inpL;
cb(h_nextn, "h_nextn", -1);
res->t_h_nextn = h_nextn;
}
if (inp_out_ids) {
ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
flat = ggml_get_rows(ctx0, flat, inp_out_ids);
inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs);
inpL = ggml_reshape_3d(ctx0, flat_out, n_embd, hc, n_outputs);
}
cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
@@ -1201,3 +1400,145 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
ggml_build_forward_expand(gf, cur);
}
llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
graph(params) {
GGML_ASSERT(hparams.n_layer_nextn > 0 && "DEEPSEEK4 MTP requires n_layer_nextn > 0");
GGML_ASSERT(hparams.n_layer_nextn == 1 && "DEEPSEEK4 MTP currently only supports a single MTP block");
GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
"nextn_layer_offset out of range [0, n_layer_nextn)");
GGML_ASSERT(ubatch.token && "DEEPSEEK4 MTP requires token input");
const int64_t hc = hparams.dsv4_hc_mult;
GGML_ASSERT(hparams.n_embd_out() == (uint32_t) (n_embd*hc) && "DEEPSEEK4 MTP hidden width mismatch");
const int il = hparams.n_layer() + cparams.nextn_layer_offset;
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd_out());
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens);
ggml_set_input(inp->embd);
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens);
ggml_set_input(inp->h);
ggml_set_name(inp->h, "mtp_h_input");
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
cb(tok_embd, "mtp_tok_embd", il);
ggml_tensor * h_state = ggml_reshape_3d(ctx0, inp->h, n_embd, hc, n_tokens);
cb(h_state, "mtp_h_state", il);
res->add_input(std::move(inp));
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
ggml_tensor * h_norm = build_norm(h_state, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
e_norm = ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens);
e_norm = ggml_repeat_4d(ctx0, e_norm, n_embd, hc, n_tokens, 1);
cb(e_norm, "mtp_enorm", il);
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0);
cb(concat, "mtp_concat", il);
ggml_tensor * inpL = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
cb(inpL, "mtp_eh_proj", 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, "mtp_hc_attn_pre", il);
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_norm", il);
cur = build_attention(model, inp_attn, cur, inp_pos, il);
inpL = build_hc_post(cur, residual, post, comb, il);
cb(inpL, "mtp_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, "mtp_hc_ffn_pre", il);
cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_ffn_norm", il);
GGML_ASSERT((uint32_t) il >= hparams.dsv4_hash_layer_count && "DEEPSEEK4 MTP does not support hash-routed MTP blocks");
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, "mtp_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, "mtp_ffn_shexp", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
cb(cur, "mtp_ffn_out", il);
inpL = build_hc_post(cur, residual, post, comb, il);
inpL = build_cvec(inpL, il);
cb(inpL, "mtp_l_out", il);
ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
ggml_tensor * h_nextn = ggml_get_rows(ctx0, flat, inp_out_ids);
cb(h_nextn, "h_nextn", -1);
res->t_h_nextn = h_nextn;
inpL = ggml_reshape_3d(ctx0, h_nextn, n_embd, hc, n_outputs);
cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);
cb(cur, "mtp_hc_head", -1);
ggml_tensor * head_norm_w = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm;
GGML_ASSERT(head_norm_w && "DEEPSEEK4 MTP missing shared head norm");
cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
cb(cur, "mtp_shared_head_norm", -1);
res->t_embd = cur;
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
GGML_ASSERT(head_w && "DEEPSEEK4 MTP missing LM head");
cur = ggml_mul_mat(ctx0, head_w, cur);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
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);
}
}
src/models/glm-dsa.cpp
+5 -1
View File
@@ -91,7 +91,11 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) {
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const bool is_mla = hparams.is_mla();
if (!is_mla) {
src/models/hy-v3.cpp
+5 -1
View File
@@ -33,7 +33,11 @@ void llama_model_hy_v3::load_arch_tensors(llama_model_loader & ml) {
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
src/models/mimo2.cpp
+5 -1
View File
@@ -30,7 +30,11 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) {
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
src/models/minimax-m3.cpp
+11 -3
View File
@@ -271,8 +271,6 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
} else {
const int64_t n_idx_dim = hparams.indexer_head_size; // 128
GGML_ASSERT(!inp_attn->self_k_rot && !inp_attn->self_v_rot && "MSA: attn-rot not supported");
// Index Branch, project, norm, partial RoPE, cache
ggml_tensor * iq = build_lora_mm(model.layers[il].index_q_proj, cur);
ggml_tensor * ik = build_lora_mm(model.layers[il].index_k_proj, cur);
@@ -289,6 +287,14 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
ggml_build_forward_expand(gf, mctx_cur->cpy_k_idx(ctx0, ik, inp_attn->get_k_idxs(), il));
ggml_tensor * ik_kv = mctx_cur->get_k_idx(ctx0, il);
if (inp_attn->self_k_rot) {
Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp_attn->self_k_rot);
Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);
}
if (inp_attn->self_v_rot) {
Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);
}
// Main branch: store K/V, take cache views
ggml_build_forward_expand(gf, Qcur);
ggml_build_forward_expand(gf, Kcur);
@@ -431,7 +437,9 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
cur = ggml_concat(ctx0, cur, outs[st], 1);
}
}
if (inp_attn->self_v_rot) {
cur = llama_mul_mat_hadamard(ctx0, cur, inp_attn->self_v_rot);
}
cb(cur, "kqv_out", il);
if (model.layers[il].wo) {
cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
src/models/models.h
+32
View File
@@ -1097,6 +1097,10 @@ struct llama_model_deepseek32 : public llama_model_base {
graph(const llama_model & model, const llm_graph_params & params);
};
struct graph_mtp : public llm_graph_context {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
@@ -1107,6 +1111,7 @@ struct llama_model_deepseek4 : public llama_model_base {
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(const llm_graph_params & params) : llm_graph_context(params) {}
graph(const llama_model & model, const llm_graph_params & params);
ggml_tensor * build_hc_pre(
@@ -1138,6 +1143,21 @@ struct llama_model_deepseek4 : public llama_model_base {
ggml_tensor * inp_pos,
int il) const;
ggml_tensor * build_attention(
const llama_model & model,
llm_graph_input_attn_k_iswa * inp_mtp,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int il) const;
ggml_tensor * build_attention_impl(
const llama_model & model,
llm_graph_input_dsv4 * inp_dsv4,
llm_graph_input_attn_k_iswa * inp_mtp,
ggml_tensor * cur,
ggml_tensor * inp_pos,
int il) const;
ggml_tensor * build_hca_compressed_kv_from_state(
ggml_tensor * kv_state,
ggml_tensor * score_state,
@@ -1213,6 +1233,10 @@ struct llama_model_deepseek4 : public llama_model_base {
int il) const;
};
struct graph_mtp : public graph {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
@@ -1272,6 +1296,10 @@ struct llama_model_dflash : public llama_model_base {
ggml_tensor * build_inp_embd_enc() const;
};
struct graph_dsv4 : public llama_model_deepseek4::graph {
graph_dsv4(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
@@ -2009,6 +2037,10 @@ struct llama_model_qwen3next : public llama_model_base {
const llama_model & model;
};
struct graph_mtp : public llm_graph_context {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
src/models/qwen35.cpp
+16 -15
View File
@@ -39,6 +39,7 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
@@ -97,25 +98,25 @@ void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {
auto & layer = layers[il];
// MTP block looks like a full-attention Qwen3.5 decoder block.
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, 0);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, 0);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, mtp_flags);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, mtp_flags);
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0);
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, mtp_flags);
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), {n_embd, n_ff}, 0);
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, mtp_flags);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd}, mtp_flags);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), {n_embd, n_ff}, mtp_flags);
// NextN-specific tensors that define the MTP block.
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, 0);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, 0);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, 0);
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, TENSOR_NOT_REQUIRED);
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags|TENSOR_NOT_REQUIRED);
};
for (int i = 0; i < n_layer; ++i) {
src/models/qwen35moe.cpp
+20 -19
View File
@@ -42,6 +42,7 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
@@ -113,32 +114,32 @@ void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff;
// MTP block looks like a full-attention Qwen3.5 decoder block with MoE FFN.
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, 0);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, 0);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, mtp_flags);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, mtp_flags);
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0);
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0);
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, mtp_flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, mtp_flags);
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, mtp_flags);
// Routed experts
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0);
create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, mtp_flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, mtp_flags);
create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, mtp_flags);
// Shared experts
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0);
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, mtp_flags);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, mtp_flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, mtp_flags);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, mtp_flags);
// NextN-specific tensors that define the MTP block.
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, 0);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, 0);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, 0);
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, TENSOR_NOT_REQUIRED);
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags|TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags|TENSOR_NOT_REQUIRED);
};
for (int i = 0; i < n_layer; ++i) {
src/models/qwen3next.cpp
+276 -49
View File
@@ -13,7 +13,11 @@ void llama_model_qwen3next::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
// Mark recurrent layers (linear attention layers)
// NextN/MTP: extra decoder block appended beyond the main stack
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
// Mark recurrent layers (linear attention layers).
if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) {
uint32_t full_attn_interval = 4;
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
@@ -28,13 +32,17 @@ void llama_model_qwen3next::load_arch_hparams(llama_model_loader & ml) {
}
}
void llama_model_qwen3next::load_arch_tensors(llama_model_loader &) {
void llama_model_qwen3next::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
if (n_expert == 0) {
throw std::runtime_error(arch_name() + " model cannot have zero experts");
}
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
// output
@@ -61,49 +69,73 @@ void llama_model_qwen3next::load_arch_tensors(llama_model_loader &) {
const int64_t qkvz_dim = key_dim * 2 + value_dim * 2;
const int64_t ba_dim = n_v_heads * 2;
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
const uint32_t n_ff_shexp = hparams.n_ff_shexp > 0 ? hparams.n_ff_shexp : hparams.n_ff(i);
auto load_block_trunk = [&](int il, int flags) {
auto & layer = layers[il];
const uint32_t n_ff_shexp = hparams.n_ff_shexp > 0 ? hparams.n_ff_shexp : hparams.n_ff(il);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);
layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", il), { n_embd }, flags);
if (!hparams.is_recr(i)) {
if (!hparams.is_recr(il)) {
// Attention layers
create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);
create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags);
// Q/K normalization for attention layers
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0);
layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, flags);
layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, flags);
} else {
// Linear attention (gated delta net) specific tensors
// Create tensors with calculated dimensions
// note: ssm_in is used by legacy GGUF
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), { n_embd, qkvz_dim }, TENSOR_NOT_REQUIRED);
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED);
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, TENSOR_NOT_REQUIRED);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, 0);
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, 0);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, 0);
layer.ssm_beta_alpha = create_tensor(tn(LLM_TENSOR_SSM_BETA_ALPHA, "weight", i), { n_embd, ba_dim }, 0);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, 0);
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, 0);
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", il), { n_embd, qkvz_dim }, TENSOR_NOT_REQUIRED | flags);
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, TENSOR_NOT_REQUIRED | flags);
layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, TENSOR_NOT_REQUIRED | flags);
layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, flags);
layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, flags);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, flags);
layer.ssm_beta_alpha = create_tensor(tn(LLM_TENSOR_SSM_BETA_ALPHA, "weight", il), { n_embd, ba_dim }, flags);
layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, flags);
layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, flags);
}
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, 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);
create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, 0);
layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags);
layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, flags);
create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, flags);
// Shared experts
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), { n_embd }, 0);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, 0);
layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, flags);
layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags);
layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags);
layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, flags);
};
auto load_block_mtp = [&](int il) {
// MTP head is identical to the trunk block (full attention + FFN)
load_block_trunk(il, mtp_flags);
auto & layer = layers[il];
// NextN-specific tensors that define the MTP block.
layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, mtp_flags);
layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, mtp_flags);
layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, mtp_flags);
layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), { n_embd, n_vocab }, mtp_flags | TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), { n_embd, n_vocab }, mtp_flags | TENSOR_NOT_REQUIRED);
layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), { n_embd }, mtp_flags | TENSOR_NOT_REQUIRED);
};
for (int i = 0; i < n_layer; i++) {
load_block_trunk(i, trunk_flags);
}
for (int i = n_layer; i < n_layer_all; i++) {
load_block_mtp(i);
}
}
std::unique_ptr<llm_graph_context> llama_model_qwen3next::build_arch_graph(const llm_graph_params & params) const {
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
return std::make_unique<graph_mtp>(*this, params);
}
return std::make_unique<graph>(*this, params);
}
@@ -120,6 +152,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
// MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.
for (int il = 0; il < n_layer; ++il) {
res->t_layer_inp[il] = inpL;
@@ -139,7 +172,7 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
cur = build_layer_attn(inp->get_attn(), cur, inp_pos, il);
}
if (il == n_layer - 1 && inp_out_ids) {
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
@@ -171,9 +204,16 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
}
cur = inpL;
// Final norm
// post-norm hidden state is input to both the LM head and the MTP head
cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cb(cur, "result_norm", -1);
res->t_embd = cur;
@@ -186,15 +226,6 @@ llama_model_qwen3next::graph::graph(const llama_model & model, const llm_graph_p
ggml_build_forward_expand(gf, cur);
}
// utility to get one slice from the third dimension
// input dim: [x, y, c, b]
// output dim: [x, y, 1, b]
// static ggml_tensor * get_slice_2d(ggml_context * ctx0, ggml_tensor * t, int64_t c) {
// return ggml_view_4d(ctx0, t, t->ne[0], t->ne[1], 1, t->ne[3],
// t->nb[1], t->nb[2], t->nb[3], t->nb[2] * c);
// }
//kcpp: already defined in delta-net-base.cpp
ggml_tensor * llama_model_qwen3next::graph::build_norm_gated(
ggml_tensor * input,
ggml_tensor * weights,
@@ -217,7 +248,7 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn(
// Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention
// Qwen3Next uses a single Q projection that outputs query + gate
ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur);
ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s);
cb(Qcur_full, "Qcur_full", il);
Qcur_full = ggml_reshape_4d(ctx0, Qcur_full, n_embd_head * 2, n_head, n_tokens, 1);
@@ -233,10 +264,10 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn(
Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full));
cb(gate, "gate", il);
ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur);
ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s);
cb(Kcur, "Kcur", il);
ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur);
ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s);
cb(Vcur, "Vcur", il);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
@@ -275,8 +306,6 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn(
gate = ggml_sigmoid(ctx0, gate);
cb(gate, "gate_sigmoid", il);
gate = ggml_reshape_2d(ctx0, gate, n_embd_head * n_head, n_tokens);
cur = ggml_mul(ctx0, cur, gate);
cb(cur, "attn_gated", il);
@@ -551,16 +580,19 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_ffn(ggml_tensor * cur, c
LLM_FFN_SILU, true,
hparams.expert_weights_scale,
LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,
nullptr, model.layers[il].ffn_gate_up_exps);
nullptr, model.layers[il].ffn_gate_up_exps,
model.layers[il].ffn_up_exps_s,
model.layers[il].ffn_gate_exps_s,
model.layers[il].ffn_down_exps_s);
cb(moe_out, "ffn_moe_out", il);
// Add shared experts if present - following Qwen3Next reference implementation
if (model.layers[il].ffn_up_shexp != nullptr) {
ggml_tensor * ffn_shexp =
build_ffn(cur,
model.layers[il].ffn_up_shexp, NULL, NULL,
model.layers[il].ffn_gate_shexp, NULL, NULL,
model.layers[il].ffn_down_shexp, NULL, NULL,
model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,
model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
NULL,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(ffn_shexp, "ffn_shexp", il);
@@ -594,3 +626,198 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_ffn(ggml_tensor * cur, c
}
return cur;
}
// LLM_GRAPH_TYPE_DECODER_MTP draft head for Qwen3-Next
llama_model_qwen3next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
: llm_graph_context(params) {
GGML_ASSERT(hparams.n_layer_nextn > 0 && "QWEN3NEXT MTP requires n_layer_nextn > 0");
GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN3NEXT MTP currently only supports a single MTP block");
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
const int il = hparams.n_layer();
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp");
// TODO: extract in a common llm_graph_context::build_inp_embd_h()
auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
ggml_set_input(inp->embd);
// TODO: make static using `ggml_build_forward_select()`
// see llm_graph_context::build_inp_embd() for reference
ggml_tensor * tok_embd;
if (ubatch.token) {
ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
} else {
tok_embd = inp->embd;
}
cb(tok_embd, "mtp_tok_embd", il);
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
ggml_set_input(inp->h);
ggml_set_name(inp->h, "mtp_h_input");
ggml_tensor * h_embd = inp->h;
res->add_input(std::move(inp));
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
cb(e_norm, "mtp_enorm", il);
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
cb(concat, "mtp_concat", il);
ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
cb(cur, "mtp_eh_proj", il);
ggml_tensor * inpSA = cur;
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_norm", il);
ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s);
cb(Qcur_full, "mtp_Qcur_full", il);
ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full,
n_embd_head, n_head, n_tokens,
ggml_element_size(Qcur_full) * n_embd_head * 2,
ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,
0);
Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il);
cb(Qcur, "mtp_Qcur_normed", il);
ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s);
Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il);
cb(Kcur, "mtp_Kcur_normed", il);
ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s);
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
cb(Qcur, "mtp_Qcur", il);
cb(Kcur, "mtp_Kcur", il);
cb(Vcur, "mtp_Vcur", il);
const float kq_scale = hparams.f_attention_scale == 0.0f
? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
cur = build_attn(inp_attn,
nullptr, nullptr, nullptr,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
cb(cur, "mtp_attn_pregate", il);
ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full,
n_embd_head, n_head, n_tokens,
ggml_element_size(Qcur_full) * n_embd_head * 2,
ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head,
ggml_element_size(Qcur_full) * n_embd_head);
// TODO: CUDA is missing non-contiguous unary ops. when implemented: remove this cont
gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens);
cb(gate, "mtp_gate", il);
cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate));
cur = build_lora_mm(layer.wo, cur, layer.wo_s);
cb(cur, "mtp_attn_out", il);
if (inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
cur = ggml_add(ctx0, cur, inpSA);
cb(cur, "mtp_attn_residual", il);
ggml_tensor * ffn_residual = cur;
cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_post_norm", il);
// MoE FFN — routed experts plus gated shared expert (mirrors the trunk).
ggml_tensor * moe_out =
build_moe_ffn(cur,
layer.ffn_gate_inp,
layer.ffn_up_exps,
layer.ffn_gate_exps,
layer.ffn_down_exps,
nullptr,
n_expert, n_expert_used,
LLM_FFN_SILU, true,
hparams.expert_weights_scale,
LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il,
nullptr, layer.ffn_gate_up_exps,
layer.ffn_up_exps_s,
layer.ffn_gate_exps_s,
layer.ffn_down_exps_s);
cb(moe_out, "mtp_ffn_moe_out", il);
if (layer.ffn_up_shexp != nullptr) {
ggml_tensor * ffn_shexp =
build_ffn(cur,
layer.ffn_up_shexp, nullptr, layer.ffn_up_shexp_s,
layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s,
layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s,
nullptr,
LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(ffn_shexp, "mtp_ffn_shexp", il);
ggml_tensor * shared_gate = build_lora_mm(layer.ffn_gate_inp_shexp, cur);
shared_gate = ggml_sigmoid(ctx0, shared_gate);
cb(shared_gate, "mtp_shared_expert_gate_sigmoid", il);
ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate);
cb(ffn_shexp, "mtp_ffn_shexp_gated", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
} else {
cur = moe_out;
}
cb(cur, "mtp_ffn_out", il);
cur = ggml_add(ctx0, cur, ffn_residual);
cb(cur, "mtp_post_ffn", il);
ggml_tensor * head_norm_w = layer.nextn.shared_head_norm
? layer.nextn.shared_head_norm
: model.output_norm;
GGML_ASSERT(head_norm_w && "QWEN3NEXT MTP: missing both nextn.shared_head_norm and output_norm");
cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
GGML_ASSERT(head_w && "QWEN3NEXT MTP: missing LM head (nextn.shared_head_head or model.output)");
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
src/models/step35.cpp
+5 -1
View File
@@ -48,7 +48,11 @@ void llama_model_step35::load_arch_tensors(llama_model_loader & ml) {
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);