GGML_ASSERT(ctx_tgt && ctx_dft && "MTP requires ctx_tgt and ctx_dft to be set");
n_embd = llama_model_n_embd_out(llama_get_model(ctx_dft));
- GGML_ASSERT(n_embd == llama_model_n_embd(llama_get_model(ctx_tgt)) &&
+ GGML_ASSERT(n_embd == llama_model_n_embd_out(llama_get_model(ctx_tgt)) &&
"MTP input row width must match the target h_nextn width");
n_mtp_layers = std::max(1, (int) llama_model_n_layer_nextn(llama_get_model(ctx_dft)));
@ModelBase.register("DeepseekV4ForCausalLM")
class DeepseekV4Model(TextModel):
model_arch = gguf.MODEL_ARCH.DEEPSEEK4
+ supports_mtp_export = True
_skipped_mtp_tensors = 0
+ _dsv4_main_layers: int | None = None
+ _dsv4_nextn_layers: int = 0
def __init__(self, *args, **kwargs):
type(self)._skipped_mtp_tensors = 0
self.hparams.setdefault(key, value)
self.block_count = self.hparams["num_hidden_layers"]
+ if self.mtp_only:
+ self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
self._dsv4_fp8_dequantized: set[str] = set()
with open(template_path, "r", encoding="utf-8") as f:
self.gguf_writer.add_chat_template(f.read())
+ def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+ type(self)._dsv4_main_layers = self.hparams["num_hidden_layers"]
+ type(self)._dsv4_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0)
+ return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
+
@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
- name, _ = item
+ name, gen = item
if name.startswith("mtp."):
- cls._skipped_mtp_tensors += 1
- return None
- return super().filter_tensors(item)
+ if not cls.mtp_only:
+ cls._skipped_mtp_tensors += 1
+ return None
+
+ assert cls._dsv4_main_layers is not None
+ parts = name.split(".", 2)
+ if len(parts) < 3 or not parts[1].isdecimal():
+ raise ValueError(f"Unexpected DeepSeek-V4 MTP tensor {name!r}")
+
+ mtp_idx = int(parts[1])
+ if mtp_idx >= cls._dsv4_nextn_layers:
+ raise ValueError(f"Unexpected DeepSeek-V4 MTP layer {mtp_idx}")
+
+ bid = cls._dsv4_main_layers + mtp_idx
+ suffix = parts[2]
+ root_hc_head = {
+ "hc_head_fn",
+ "hc_head_base",
+ "hc_head_scale",
+ }
+ if suffix in root_hc_head:
+ name = suffix
+ elif suffix in (
+ "e_proj.weight", "e_proj.scale",
+ "h_proj.weight", "h_proj.scale",
+ ):
+ name = f"layers.{bid}.nextn.{suffix}"
+ elif suffix == "enorm.weight":
+ name = f"layers.{bid}.nextn.enorm.weight"
+ elif suffix == "hnorm.weight":
+ name = f"layers.{bid}.nextn.hnorm.weight"
+ elif suffix == "norm.weight":
+ name = f"layers.{bid}.nextn.shared_head_norm.weight"
+ else:
+ name = f"layers.{bid}.{suffix}"
+ return name, gen
+
+ if cls.mtp_only:
+ keep = name in (
+ "embed.weight",
+ "norm.weight",
+ "head.weight",
+ "head.scale",
+ )
+ if not keep:
+ return None
+
+ return super().filter_tensors((name, gen))
@staticmethod
def _float8_dtypes() -> tuple[torch.dtype, ...]:
self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hparams["hc_sinkhorn_iters"])
self.gguf_writer.add_hyper_connection_epsilon(hparams["hc_eps"])
self.gguf_writer.add_hash_layer_count(hparams["num_hash_layers"])
+ self.gguf_writer.add_embedding_length_out(hparams["hidden_size"] * hparams["hc_mult"])
+ if self.mtp_only and (num_nextn_predict_layers := hparams.get("num_nextn_predict_layers", 0)) > 0:
+ self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
def dequant_model(self):
fp8_dtypes = self._float8_dtypes()
if self._dsv4_mxfp4_generated:
return ()
- consumed: list[str] = self._write_hash_routing_tensors()
+ consumed: list[str] = []
+ main_layers = self.hparams["num_hidden_layers"]
+ if not self.mtp_only:
+ consumed.extend(self._write_hash_routing_tensors())
+ elif self.hparams["num_hash_layers"] > 0:
+ for bid in range(self.hparams["num_hash_layers"]):
+ name = f"layers.{bid}.ffn.gate.tid2eid"
+ if name in self.model_tensors:
+ consumed.extend(self._write_hash_routing_tensors())
+ break
+
for bid in range(self.block_count):
+ if self.mtp_only and bid < main_layers:
+ continue
consumed.extend(self._write_mxfp4_expert_tensor(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP))
consumed.extend(self._write_mxfp4_expert_tensor(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP))
consumed.extend(self._write_mxfp4_expert_tensor(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP))
+ for bid in range(main_layers, self.block_count):
+ e_name = f"layers.{bid}.nextn.e_proj.weight"
+ h_name = f"layers.{bid}.nextn.h_proj.weight"
+ if e_name not in self.model_tensors and h_name not in self.model_tensors:
+ continue
+ if e_name not in self.model_tensors or h_name not in self.model_tensors:
+ raise KeyError(f"Missing DeepSeek-V4 MTP e/h projection pair for block {bid}")
+
+ e_proj = LazyTorchTensor.to_eager(self.model_tensors[e_name]())
+ h_proj = LazyTorchTensor.to_eager(self.model_tensors[h_name]())
+ yield (f"layers.{bid}.nextn.eh_proj.weight", torch.cat((e_proj, h_proj), dim=1).contiguous())
+ consumed.extend((e_name, h_name))
+
for name in consumed:
del self.model_tensors[name]
"ffn.shared_experts.w1.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"),
"ffn.shared_experts.w2.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"),
"ffn.shared_experts.w3.weight": (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"),
+ "nextn.eh_proj.weight": (gguf.MODEL_TENSOR.NEXTN_EH_PROJ, ".weight"),
+ "nextn.enorm.weight": (gguf.MODEL_TENSOR.NEXTN_ENORM, ".weight"),
+ "nextn.hnorm.weight": (gguf.MODEL_TENSOR.NEXTN_HNORM, ".weight"),
+ "nextn.shared_head_norm.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ".weight"),
+ "nextn.embed_tokens.weight": (gguf.MODEL_TENSOR.NEXTN_EMBED_TOKENS, ".weight"),
+ "nextn.shared_head_head.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, ".weight"),
}
tensor_name = match.group(2)
return [(self._format_dsv4_tensor_name(tensor_key, bid, suffix), data_torch)]
def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
- del new_name, bid # unused
+ del bid # unused
if name in self._dsv4_fp8_dequantized and n_dims >= 2:
return gguf.GGMLQuantizationType.Q8_0
+ if new_name.endswith(".nextn.eh_proj.weight"):
+ return gguf.GGMLQuantizationType.Q8_0
if name in self._dsv4_f32_tensors:
return gguf.GGMLQuantizationType.F32
if name in self._dsv4_bf16_tensors and n_dims >= 2:
return False
+ def prepare_metadata(self, vocab_only: bool):
+ from_dir = self.fname_out.is_dir()
+ super().prepare_metadata(vocab_only=vocab_only)
+
+ if not self.mtp_only or not from_dir:
+ return
+
+ output_type: str = self.ftype.name.partition("_")[2]
+ fname_default: str = gguf.naming_convention(
+ self.metadata.name, self.metadata.basename, self.metadata.finetune,
+ self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+ self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
def prepare_tensors(self):
super().prepare_tensors()
self._is_mxfp4 = True
MODEL_TENSOR.FFN_GATE_SHEXP,
MODEL_TENSOR.FFN_DOWN_SHEXP,
MODEL_TENSOR.FFN_UP_SHEXP,
+ MODEL_TENSOR.NEXTN_EH_PROJ,
+ MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+ MODEL_TENSOR.NEXTN_ENORM,
+ MODEL_TENSOR.NEXTN_HNORM,
+ MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+ MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
],
MODEL_ARCH.ERNIE4_5_MOE: [
MODEL_TENSOR.TOKEN_EMBD,
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
+ case LLM_ARCH_DEEPSEEK4:
return true;
default:
return false;
switch (arch) {
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
+ case LLM_ARCH_DEEPSEEK4:
return true;
default:
return false;
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;
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;
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;
}
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;
}
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]);
}
}
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());
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;
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);
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());
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");
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();
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);
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);
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);
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();
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 {
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]
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]
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(
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;
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);
}
}
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);
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) {
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;
}
}
- 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);
}
}
}
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;
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;
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);
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);
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 {
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);
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) {
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;
}
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);
}
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) {
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();
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);
}
}
}
}
-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(ggml_context * ctx, int32_t il) const {
+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_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_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 {
+ ggml_tensor * state = get_score_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::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 {
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)) {
// 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);
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
}
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;
+ }
+
+ 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);
+
+ return res;
+ }
+
+ if (n_rs_seq == 0) {
return false;
}
- bool res = true;
+ const llama_pos rollback = pos_max - (p0 - 1);
+ if (rollback < 1 || rollback > (llama_pos) n_rs_seq) {
+ return false;
+ }
- 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);
+ const bool res = kv_raw->seq_rm(seq_id, p0, p1);
+ if (res) {
+ rs_idx[seq_id] = (uint32_t) rollback;
+ }
return res;
}
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) {
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) {
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 {
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);
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);
+ }
}
//
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),
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;
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;
}
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;
}
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;
}
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);
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;
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;
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);
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;
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;
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
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:
}
}
- 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) {
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;
}
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);
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;
}
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);
}
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;
}
}
-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;
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);
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];
-
- 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);
+ const int flags = i < n_layer ? trunk_flags : mtp_flags;
+
+ 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}, 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}, flags);
+
+ 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}, 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);
}
- 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);
}
}
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);
}
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 * 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 * zero = ggml_scale(ctx, ggml_sum(ctx, dep), 0.0f);
- return ggml_add(ctx, t, zero);
+ 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,
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();
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);
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,
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);
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,
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);
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,
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);
}
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) {
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;
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);
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);
+}
}
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) {
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];
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");
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);
+ }
+}
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(
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,
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;
};
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;
};
if (llama_vocab_type(vocab) == LLAMA_VOCAB_TYPE_NONE) {
tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 };
} else {
- tokens = common_tokenize(ctx_src, "The quick brown fox jumps", true);
+ tokens = common_tokenize(ctx_src, "The quick brown fox jumps over the lazy dog", true);
}
const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src);
- if (tokens.size() > n_rs_seq + 1) {
- tokens.resize(n_rs_seq + 1);
+ constexpr uint32_t n_rollback = 3;
+ if (n_rs_seq < n_rollback) {
+ fprintf(stderr, "%s : skipping because n_rs_seq is too small\n", __func__);
+ llama_free(ctx_src);
+ llama_free(ctx_dst);
+ return 0;
}
- if (tokens.size() < 2) {
+ if (tokens.empty()) {
fprintf(stderr, "%s : not enough prompt tokens\n", __func__);
return 1;
}
- const uint32_t n_tokens = tokens.size();
- const llama_token last_tok = tokens.back();
- const llama_pos last_pos = (llama_pos) n_tokens - 2;
+ tokens.resize(n_rs_seq + 1, tokens.back());
- // Decode the full prompt on the source, then roll back the last position.
+ const uint32_t n_tokens = tokens.size();
+ const llama_pos rollback_pos = (llama_pos) n_tokens - n_rollback;
+
+ // Decode the full prompt on the source, then roll back three positions.
+ // Replaying them crosses DSV4's ratio-4 compressor boundary.
// Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0).
if (!decode_tokens(ctx_src, tokens, n_tokens)) {
fprintf(stderr, "%s : failed to decode prompt\n", __func__);
return 1;
}
- if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, last_pos, -1)) {
+ if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1)) {
fprintf(stderr, "%s : rollback failed\n", __func__);
return 1;
}
ckpt.update_tgt(ctx_src, 0, 0);
ckpt.load_tgt(ctx_dst, 0, 0);
- // Replay the rolled-back token on both contexts and compare logits.
- if (!decode_one(ctx_src, last_tok, last_pos) ||
- !decode_one(ctx_dst, last_tok, last_pos)) {
- fprintf(stderr, "%s : replay failed\n", __func__);
+ constexpr float eps = 1e-5f;
+ std::vector<std::vector<float>> logits_src_replay(n_rollback);
+ const auto replay_and_compare = [&](const char * mode) {
+ for (uint32_t i = 0; i < n_rollback; ++i) {
+ const llama_pos pos = rollback_pos + i;
+ if (!decode_one(ctx_src, tokens[pos], pos) ||
+ !decode_one(ctx_dst, tokens[pos], pos)) {
+ fprintf(stderr, "%s : %s replay failed at position %d\n", __func__, mode, pos);
+ return false;
+ }
+
+ const float * logits_src = llama_get_logits_ith(ctx_src, 0);
+ const float * logits_dst = llama_get_logits_ith(ctx_dst, 0);
+ if (logits_src == nullptr || logits_dst == nullptr) {
+ fprintf(stderr, "%s : missing %s logits at position %d\n", __func__, mode, pos);
+ return false;
+ }
+
+ logits_src_replay[i].assign(logits_src, logits_src + n_vocab);
+ for (int token = 0; token < n_vocab; ++token) {
+ if (std::fabs(logits_src[token] - logits_dst[token]) > eps) {
+ fprintf(stderr, "%s : %s logits mismatch at position %d, token %d (%g != %g)\n",
+ __func__, mode, pos, token, (double) logits_src[token], (double) logits_dst[token]);
+ return false;
+ }
+ }
+ }
+ return true;
+ };
+ if (!replay_and_compare("full")) {
return 1;
}
- const float * logits_src = llama_get_logits_ith(ctx_src, 0);
- const float * logits_dst = llama_get_logits_ith(ctx_dst, 0);
- if (logits_src == nullptr || logits_dst == nullptr) {
- fprintf(stderr, "%s : missing logits\n", __func__);
+ if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) ||
+ !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) {
+ fprintf(stderr, "%s : partial rollback failed\n", __func__);
return 1;
}
- constexpr float eps = 1e-5f;
- for (int i = 0; i < n_vocab; ++i) {
- if (std::fabs(logits_src[i] - logits_dst[i]) > eps) {
- fprintf(stderr, "%s : logits mismatch at token %d (%g != %g)\n",
- __func__, i, (double) logits_src[i], (double) logits_dst[i]);
- return 1;
- }
+ constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY;
+ common_prompt_checkpoint ckpt_partial;
+ ckpt_partial.update_tgt(ctx_src, 0, partial_flags);
+ ckpt_partial.load_tgt(ctx_dst, 0, partial_flags);
+
+ if (!replay_and_compare("partial")) {
+ return 1;
}
// Repeat the load into a context that already has its own rollback state:
- // groups 1..n_rs_seq hold a *different* prompt's history, and rs_idx[0] is
+ // groups 1..n_rs_seq hold a different prompt's history, and rs_idx[0] is
// non-zero at load time. The restore must wipe that state and still match.
llama_context * ctx_dirty = make_ctx(params, model);
if (ctx_dirty == nullptr) {
fprintf(stderr, "%s : dirty prompt decode failed\n", __func__);
return 1;
}
- if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, last_pos, -1)) {
+ if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, rollback_pos, -1)) {
fprintf(stderr, "%s : dirty rollback failed\n", __func__);
return 1;
}
ckpt.load_tgt(ctx_dirty, 0, 0);
- if (!decode_one(ctx_dirty, last_tok, last_pos)) {
- fprintf(stderr, "%s : dirty replay failed\n", __func__);
- return 1;
- }
-
- const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0);
- if (logits_dirty == nullptr) {
- fprintf(stderr, "%s : missing dirty logits\n", __func__);
- return 1;
- }
+ for (uint32_t i = 0; i < n_rollback; ++i) {
+ const llama_pos pos = rollback_pos + i;
+ if (!decode_one(ctx_dirty, tokens[pos], pos)) {
+ fprintf(stderr, "%s : dirty replay failed at position %d\n", __func__, pos);
+ return 1;
+ }
- for (int i = 0; i < n_vocab; ++i) {
- if (std::fabs(logits_src[i] - logits_dirty[i]) > eps) {
- fprintf(stderr, "%s : dirty-ctx logits mismatch at token %d (%g != %g)\n",
- __func__, i, (double) logits_src[i], (double) logits_dirty[i]);
+ const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0);
+ if (logits_dirty == nullptr) {
+ fprintf(stderr, "%s : missing dirty logits at position %d\n", __func__, pos);
return 1;
}
+
+ for (int token = 0; token < n_vocab; ++token) {
+ if (std::fabs(logits_src_replay[i][token] - logits_dirty[token]) > eps) {
+ fprintf(stderr, "%s : dirty-ctx logits mismatch at position %d, token %d (%g != %g)\n",
+ __func__, pos, token, (double) logits_src_replay[i][token], (double) logits_dirty[token]);
+ return 1;
+ }
+ }
}
fprintf(stderr, "%s : recurrent rollback checkpoint restored successfully\n", __func__);
overview / pkg / ggml / sources / llama.cpp / commitdiff