diff --git a/common/chat.cpp b/common/chat.cpp index 99ed74faf7a9..d9c2bcc3370f 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -428,6 +428,41 @@ std::vector common_chat_msgs_parse_oaicompat(const json & messa if (message.contains("tool_call_id")) { msg.tool_call_id = message.at("tool_call_id"); } + { + json extra = json::object(); + for (const auto & [k, v] : message.items()) { + if (k != "role" && k != "content" && k != "tool_calls" && k != "reasoning_content" && + k != "name" && k != "tool_call_id") { + extra[k] = v; + } + } + if (has_tool_calls) { + // per tool call: keys besides type/function/id, and function keys besides name/arguments + json tc_extra = json::array(); + bool any = false; + for (const auto & tool_call : message.at("tool_calls")) { + json e = json::object(); + for (const auto & [k, v] : tool_call.items()) { + if (k != "type" && k != "function" && k != "id") { + e[k] = v; + } + } + for (const auto & [k, v] : tool_call.at("function").items()) { + if (k != "name" && k != "arguments") { + e["function"][k] = v; + } + } + any = any || !e.empty(); + tc_extra.push_back(e); + } + if (any) { + extra["__tool_call_extras"] = tc_extra; + } + } + if (!extra.empty()) { + msg.extra_json = extra.dump(); + } + } msgs.push_back(msg); } @@ -516,6 +551,51 @@ static json render_message_to_json(const std::vector & msgs, co return messages_inp_normalizer(c).normalize(messages); } +// DeepSeek V4.1: its template reads non-OpenAI message keys ("task", per-message "tools", tool call +// "namespace"), and joins content blocks with "\n\n" around images, so keep media markers as typed text +// parts instead of the default glued string and restore the extra keys. +static json render_message_to_json_deepseek_v4_1(const std::vector & msgs, const jinja::caps & c) { + json messages = json::array(); + for (const auto & msg : msgs) { + json jmsg; + if (msg.contains_media() && msg.content.empty()) { + common_chat_msg typed = msg; + for (auto & part : typed.content_parts) { + part.type = "text"; + } + jmsg = typed.to_json_oaicompat(/* concat_typed_text= */ false); + } else { + jmsg = msg.to_json_oaicompat(/* concat_typed_text= */ false); + } + if (!msg.extra_json.empty()) { + const json extra = json::parse(msg.extra_json); + for (const auto & [k, v] : extra.items()) { + if (k == "__tool_call_extras") { + if (!jmsg.contains("tool_calls")) { + continue; + } + auto & tcs = jmsg.at("tool_calls"); + for (size_t i = 0; i < tcs.size() && i < v.size(); i++) { + for (const auto & [tk, tv] : v[i].items()) { + if (tk == "function") { + for (const auto & [fk, fv] : tv.items()) { + tcs[i]["function"][fk] = fv; + } + } else { + tcs[i][tk] = tv; + } + } + } + } else if (!jmsg.contains(k)) { + jmsg[k] = v; + } + } + } + messages.push_back(std::move(jmsg)); + } + return messages_inp_normalizer(c).normalize(messages); +} + // DEPRECATED: only used in tests json common_chat_msgs_to_json_oaicompat(const std::vector & msgs, bool concat_typed_text) { jinja::caps c; @@ -1208,11 +1288,12 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const // DeepSeek V3.2/V4 format detection: template defines dsml_token and uses it for tool calls. // The template source contains the token as a variable assignment, not as a literal in markup. // V3.2 names the tool call block "function_calls", V4 names it "tool_calls". - if (src.find("dsml_token") != std::string::npos && - src.find("DSML") != std::string::npos && + // V4.1 names it " calls" (leading-space tag names) and its template may not name the dsml_token variable. + if (((src.find("dsml_token") != std::string::npos && src.find("DSML") != std::string::npos) || + src.find("｜DSML｜ calls>") != std::string::npos) && (src.find("function_calls") != std::string::npos || src.find("tool_calls") != std::string::npos)) { - LOG_DBG("Using specialized template: DeepSeek V3.2/V4\n"); + LOG_DBG("Using specialized template: DeepSeek V3.2/V4/V4.1\n"); return common_chat_params_init_deepseek_v3_2; } @@ -1313,7 +1394,9 @@ static common_chat_params common_chat_templates_apply_jinja(const struct common_ workaround::trim_all_content(trimmed_messages); messages_to_render = &trimmed_messages; } - params.messages = render_message_to_json(*messages_to_render, tmpl.original_caps()); + params.messages = src.find("｜DSML｜ calls>") != std::string::npos ? + render_message_to_json_deepseek_v4_1(*messages_to_render, tmpl.original_caps()) : + render_message_to_json(*messages_to_render, tmpl.original_caps()); params.tool_choice = inputs.tool_choice; params.reasoning_format = inputs.reasoning_format; params.enable_thinking = inputs.enable_thinking; diff --git a/common/chat.h b/common/chat.h index dcff94919819..4569f8d9feef 100644 --- a/common/chat.h +++ b/common/chat.h @@ -94,6 +94,9 @@ struct common_chat_msg { std::string reasoning_content; std::string tool_name; std::string tool_call_id; + // Message keys the OpenAI fields above do not cover (e.g. "task", per-message "tools", tool call + // "namespace"), kept as a JSON object string. Only templates that read them get them (DeepSeek V4.1). + std::string extra_json; common_json to_json_oaicompat(bool concat_typed_text = false) const; diff --git a/common/parsers/deepseek.cpp b/common/parsers/deepseek.cpp index 7ee38ce6df74..93b5d5270be5 100644 --- a/common/parsers/deepseek.cpp +++ b/common/parsers/deepseek.cpp @@ -1,4 +1,5 @@ #include "parsers.h" +#include "log.h" // The DeepSeek V4 reference implementation renders consecutive tool results into a single // user block, ordered by the tool call order of the preceding assistant message (matched @@ -75,13 +76,31 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ // non-thinking generation prompt ends with a bare instead of an empty // pair. const bool is_v4 = tmpl.source().find("function_calls") == std::string::npos; + // V4.1 keeps the V4 behaviour but its DSML tag names carry a leading space: + // <｜DSML｜ calls>, <｜DSML｜ invoke name="...">, <｜DSML｜ parameter name="..." string="...">. + const bool is_v41 = is_v4 && tmpl.source().find("｜DSML｜ calls>") != std::string::npos; std::optional adjusted_messages; if (is_v4) { adjusted_messages = deepseek_v4_sort_tool_results(inputs.messages); } - auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + // V4.1 can declare tools on a system message instead of top-level "tools": parse calls to those too. + json tools = inputs.tools; + if (is_v41 && !(tools.is_array() && !tools.empty()) && inputs.messages.is_array()) { + for (const auto & msg : inputs.messages) { + if (msg.value("role", "") == "system" && msg.contains("tools") && msg.at("tools").is_array()) { + if (!tools.is_array()) { + tools = json::array(); + } + for (const auto & tool : msg.at("tools")) { + tools.push_back(tool); + } + } + } + } + + auto has_tools = tools.is_array() && !tools.empty(); auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); @@ -90,17 +109,40 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ if (is_v4 && has_response_format) { additional_context = json{ { "response_format", inputs.json_schema } }; } + // V4.1 accepts low/high/max or 1-100 and raises on anything else: map the remaining OpenAI values + // (minimal -> low, xhigh -> max, medium -> unset = template default) instead of failing the request. + if (is_v41 && inputs.extra_context.is_object() && inputs.extra_context.contains("reasoning_effort") && + inputs.extra_context.at("reasoning_effort").is_string()) { + const std::string effort = inputs.extra_context.at("reasoning_effort").get(); + json mapped; + if (effort == "minimal") { + mapped = "low"; + } else if (effort == "xhigh") { + mapped = "max"; + } else if (effort == "medium") { + mapped = nullptr; + } + if (effort == "minimal" || effort == "xhigh" || effort == "medium") { + LOG_DBG("%s: DeepSeek V4.1 reasoning_effort '%s' -> %s\n", __func__, effort.c_str(), + mapped.is_null() ? "unset (template default)" : mapped.get().c_str()); + if (!additional_context.has_value()) { + additional_context = json::object(); + } + (*additional_context)["reasoning_effort"] = mapped; + } + } const std::string DSML = "｜DSML｜"; const std::string THINK_START = ""; const std::string THINK_END = ""; - const std::string TC_BLOCK = is_v4 ? "tool_calls" : "function_calls"; - const std::string FC_START = "<" + DSML + TC_BLOCK + ">"; - const std::string FC_END = ""; - const std::string INVOKE_START = "<" + DSML + "invoke"; - const std::string INVOKE_END = ""; - const std::string PARAM_START = "<" + DSML + "parameter"; - const std::string PARAM_END = ""; + const std::string TAG_SP = is_v41 ? " " : ""; + const std::string TC_BLOCK = is_v41 ? "calls" : is_v4 ? "tool_calls" : "function_calls"; + const std::string FC_START = "<" + DSML + TAG_SP + TC_BLOCK + ">"; + const std::string FC_END = ""; + const std::string INVOKE_START = "<" + DSML + TAG_SP + "invoke"; + const std::string INVOKE_END = ""; + const std::string PARAM_START = "<" + DSML + TAG_SP + "parameter"; + const std::string PARAM_END = ""; const std::string GEN_PROMPT = "<｜Assistant｜>"; const std::string TC_SEPARATOR = "\n\n"; @@ -152,7 +194,7 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ // build tool call section first since we might need it in reasoning auto tool_choice = p.choice(); if (has_tool_calls) { - foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) { + foreach_function(tools, [&](size_t tool_index, const json & tool) { const auto & function = tool.at("function"); std::string name = function.at("name"); @@ -223,7 +265,11 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ } if (extract_reasoning && inputs.enable_thinking) { - reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); + // V4 renders a content continuation without reasoning as "<｜Assistant｜>" + content even in + // thinking mode (generation_prompt above), so the bare THINK_END must be consumed here as well, or it + // leaks into the final message content (the streamed deltas, which start after the prefill, omit it) + reasoning = is_v4 ? p.optional((THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END) | p.literal(THINK_END)) + : p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END); reasoning_with_tc = THINK_START + p.reasoning(p.until_one_of({ TC_SEPARATOR + FC_START, FC_START, THINK_END })) + p.space() + obligatory_tool_calls; diff --git a/conversion/__init__.py b/conversion/__init__.py index e3c10d5fc54e..81d87c836e87 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -68,6 +68,7 @@ "LingDSparkModel": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DeepseekV4DSparkModel": "deepseek", + "DeepseekV41ForCausalLM": "deepseek41", "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", "DistilBertModel": "bert", diff --git a/conversion/deepseek41.py b/conversion/deepseek41.py new file mode 100644 index 000000000000..2d2b19b245ae --- /dev/null +++ b/conversion/deepseek41.py @@ -0,0 +1,652 @@ +from __future__ import annotations + +import os +import re +from pathlib import Path + +from typing import Any, Callable, Iterable, TYPE_CHECKING + +import numpy as np +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger +from .deepseek import DeepseekV4Model + + +# --------------------------------------------------------------------------------------------------------------------- +# Engram hash constants, ported from inference/engram.py of deepseek-ai/DeepSeek-V4.1-Flash-1017. +# They are written to GGUF so that the runtime never recomputes them (the compressed token map needs the HF +# tokenizers normalizers, the multipliers numpy's PCG64 stream). +# --------------------------------------------------------------------------------------------------------------------- + +def _is_prime(n: int) -> bool: + # deterministic Miller-Rabin for n < 3.3e24 + if n < 2: + return False + small = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41) + for p in small: + if n % p == 0: + return n == p + d, s = n - 1, 0 + while d % 2 == 0: + d //= 2 + s += 1 + for a in small: + x = pow(a, d, n) + if x in (1, n - 1): + continue + for _ in range(s - 1): + x = x * x % n + if x == n - 1: + break + else: + return False + return True + + +def engram_primes(n_layers: int, max_ngram_size: int, n_heads: int, vocab_size: int) -> list[list[list[int]]]: + """EngramLayout.from_args: [layer][ngram order 2..max][head] bucket moduli, never reused across layers.""" + seen: set[int] = set() + out: list[list[list[int]]] = [] + for _ in range(n_layers): + per_ngram = [] + for _ in range(max_ngram_size - 1): + sizes, cur = [], vocab_size - 1 + for _ in range(n_heads): + cand = cur + 1 + while not _is_prime(cand) or cand in seen: + cand += 1 + cur = cand + seen.add(cur) + sizes.append(cur) + per_ngram.append(sizes) + out.append(per_ngram) + return out + + +def engram_multipliers(layer_ids: list[int], max_ngram_size: int, compressed_vocab_size: int) -> list[list[int]]: + """compute_hash_multipliers: odd, bounded so that token_id * multiplier cannot overflow int64.""" + max_long = np.iinfo(np.int64).max + bound = max(1, (max_long // compressed_vocab_size) // 2) + rows = [] + for layer_id in layer_ids: + rng = np.random.default_rng(10007 * layer_id) + values = rng.integers(low=0, high=bound, size=(max_ngram_size,), dtype=np.int64) + rows.append([int(v) * 2 + 1 for v in values]) + return rows + + +def engram_compressed_token_map(tokenizer_dir: Path) -> tuple[list[int], int]: + """build_compressed_token_map: token id -> compressed id (tokens that normalize alike collapse).""" + from tokenizers import Regex, normalizers + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir) + sentinel = "\ue000" + normalizer = normalizers.Sequence([ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), sentinel), + normalizers.Strip(), + normalizers.Replace(sentinel, " "), + ]) + backend = tokenizer.backend_tokenizer + key_to_new: dict[str, int] = {} + lookup = [0] * len(tokenizer) + for token_id in range(len(tokenizer)): + text = backend.decode([token_id], skip_special_tokens=False) + if "\ufffd" in text: + key = backend.id_to_token(token_id) + else: + normalized = normalizer.normalize_str(text) + key = normalized if normalized else text + new_id = key_to_new.get(key) + if new_id is None: + new_id = len(key_to_new) + key_to_new[key] = new_id + lookup[token_id] = new_id + return lookup, len(key_to_new) + + +# --------------------------------------------------------------------------------------------------------------------- +# Streaming raw tensors: gguf_writer keeps every added tensor until the final write and only needs +# shape / dtype / nbytes / tofile, so producing the bytes chunk by chunk at write time keeps memory flat +# (an Engram table is ~101 GB, a stacked MXFP4 expert tensor ~2.4 GB). +# --------------------------------------------------------------------------------------------------------------------- + +class _StreamedBytes: + def __init__(self, shape: tuple[int, ...], chunks: Callable[[], Iterable[np.ndarray]]): + self.shape = shape + self.dtype = np.dtype(np.uint8) + self.nbytes = int(np.prod(shape)) + self._chunks = chunks + + def tofile(self, fout) -> None: + written = 0 + for chunk in self._chunks(): + chunk = np.ascontiguousarray(chunk, dtype=np.uint8) + chunk.tofile(fout) + written += chunk.nbytes + if written != self.nbytes: + raise ValueError(f"streamed {written} bytes, expected {self.nbytes}") + + def byteswap(self, inplace: bool = False): + return self # raw bytes + + +def _e8m0_check(scale_u8: np.ndarray, name: str, lo: int, hi: int) -> None: + if scale_u8.size == 0: + return + smin, smax = int(scale_u8.min()), int(scale_u8.max()) + if smin < lo or smax > hi: + raise ValueError(f"{name}: E8M0 scale exponents [{smin}, {smax}] outside the exact range [{lo}, {hi}] " + f"(0xFF is NaN, very small exponents underflow BF16)") + + +@ModelBase.register("DeepseekV41ForCausalLM") +@ModelBase.example("deepseek-ai/DeepSeek-V4.1-Flash-1017") +class DeepseekV41Model(DeepseekV4Model): + """DeepSeek-V4.1 text model (arch deepseek41), see TENSORS.md in the port branch. + + Lossless reference layout: FP8 32x32-block linears dequantised to BF16 (exact: e4m3 * 2^k fits BF16), routed + experts copied bit-exact as MXFP4, Engram tables copied bit-exact as MXFP8. DSpark blocks (mtp.0..2) go to a + separate GGUF with --mtp. The vision tower is not converted yet. + + Environment overrides: + DSV41_DENSE_TYPE bf16 | q8_0 for the 2D dense tensors (default: q8_0 with --outtype q8_0, else bf16) + DSV41_ENGRAM_TYPE mxfp8 (default, lossless) | bf16 (lossless, 2x larger) | f32 + """ + model_arch = gguf.MODEL_ARCH.DEEPSEEK41 + supports_mtp_export = True + + # tensors that stay in their source float dtype (bf16 in the checkpoint, never quantised) + _ENGRAM_EMBED_RE = re.compile(r"layers\.(\d+)\.engram\.embed\.(weight|scale)$") + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + hp = self.hparams + # dense (non-expert, non-Engram) linears follow --outtype: bf16 is lossless, q8_0 is the B2 Q8_0 companion + default_dense = {gguf.LlamaFileType.MOSTLY_Q8_0: "q8_0", gguf.LlamaFileType.ALL_F32: "f32"}.get(self.ftype, "bf16") + self._dense_type = os.environ.get("DSV41_DENSE_TYPE", default_dense).lower() + self._engram_type = os.environ.get("DSV41_ENGRAM_TYPE", "mxfp8").lower() + if self._dense_type not in ("bf16", "q8_0", "f32"): + raise ValueError(f"DSV41_DENSE_TYPE={self._dense_type!r}") + if self._engram_type not in ("mxfp8", "bf16", "f32"): + raise ValueError(f"DSV41_ENGRAM_TYPE={self._engram_type!r}") + + qc = hp.get("quantization_config") or {} + self._fp8_block = tuple(qc.get("weight_block_size") or (128, 128)) + if qc.get("scale_fmt", "ue8m0") != "ue8m0": + raise ValueError(f"unsupported FP8 scale format {qc.get('scale_fmt')!r}") + + # the V4 constructor installed the V4 template; V4.1 has its own + template_path = Path(__file__).parent.parent / "models" / "templates" / "deepseek-ai-DeepSeek-V4.1-Flash-1017.jinja" + if template_path.is_file(): + self.gguf_writer.add_chat_template(template_path.read_text(encoding="utf-8")) + + # ---- tensor selection ------------------------------------------------------------------------------------------- + + def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]: + hp = {**self.hparams, **self.hparams.get("text_config", {})} # text_config is merged only after indexing + type(self)._dsv4_main_layers = hp["num_hidden_layers"] + type(self)._dsv4_nextn_layers = hp.get("num_nextn_predict_layers", 0) + return TextModel.index_tensors(self, remote_hf_model_id=remote_hf_model_id) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if name.startswith(("vision.", "aligner.", "image_")): + return None # vision tower: mmproj not converted yet + if name.startswith("mtp."): + if not cls.mtp_only: + cls._skipped_mtp_tensors += 1 + return None + m = re.match(r"mtp\.(\d+)\.(.+)$", name) + if m is None: + raise ValueError(f"Unexpected DeepSeek-V4.1 DSpark tensor {name!r}") + stage, rest = int(m.group(1)), m.group(2) + assert cls._dsv4_main_layers is not None + if stage >= cls._dsv4_nextn_layers: + raise ValueError(f"Unexpected DSpark stage {stage}") + bid = cls._dsv4_main_layers + stage + if rest.startswith(("markov_head.", "confidence_head.")): + return f"dspark.{rest}", gen + return f"layers.{bid}.{rest}", gen + if cls.mtp_only and name not in ("embed.weight", "norm.weight", "head.weight"): + return None + return TextModel.filter_tensors((name, gen)) + + # ---- hparams ---------------------------------------------------------------------------------------------------- + + def set_gguf_parameters(self): + TextModel.set_gguf_parameters(self) + hp = self.hparams + w = self.gguf_writer + n_main = hp["num_hidden_layers"] + + w.add_rope_dimension_count(hp["qk_rope_head_dim"]) + w.add_q_lora_rank(hp["q_lora_rank"]) + w.add_sliding_window(hp["sliding_window"]) + w.add_expert_feed_forward_length(hp["moe_intermediate_size"]) + w.add_expert_shared_count(hp["n_shared_experts"]) + w.add_expert_weights_scale(hp["routed_scaling_factor"]) + w.add_expert_weights_norm(hp["norm_topk_prob"]) + w.add_swiglu_clamp_exp([hp["swiglu_limit"]] * self.block_count) + w.add_swiglu_clamp_shexp([hp["swiglu_limit"]] * self.block_count) + + w.add_indexer_head_count(hp["index_n_heads"]) + w.add_indexer_key_length(hp["index_head_dim"]) + w.add_indexer_top_k(hp["index_topk"]) + + w.add_attention_output_group_count(hp["o_groups"]) + w.add_attention_output_lora_rank(hp["o_lora_rank"]) + w.add_attention_compress_ratios(hp["compress_ratios"]) + w.add_attention_compress_rope_freq_base(hp["compress_rope_theta"]) + w.add_hyper_connection_count(hp["hc_mult"]) + w.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"]) + w.add_hyper_connection_epsilon(hp["hc_eps"]) + w.add_embedding_length_out(hp["hidden_size"] * hp["hc_mult"]) + + # cross-layer sharing and candidate blocks + w.add_attention_kv_source_layers(hp["kv_source_layer_ids"]) + w.add_indexer_source_layers(hp["index_source_layer_ids"]) + w.add_indexer_candidate_source_layer(int(hp.get("candidate_source_layer_id", -1))) + w.add_indexer_candidate_top_k_blocks(hp["candidate_topk_blocks"]) + w.add_indexer_candidate_block_size(hp["candidate_block_size"]) + + # Engram (keys omitted when the model has no Engram layers: GGUF cannot hold empty arrays, the loader treats them as optional) + layer_ids = list(hp.get("engram_layer_ids") or []) + if layer_ids: + max_ngram, n_heads = hp["engram_max_ngram_size"], hp["engram_n_heads"] + primes = engram_primes(len(layer_ids), max_ngram, n_heads, hp["engram_vocab_size"]) + for li, layer_primes in enumerate(primes): + total = sum(p for per in layer_primes for p in per) + if total != hp["engram_num_embeddings"][li]: + raise ValueError(f"engram layer {layer_ids[li]}: prime sum {total} != num_embeddings {hp['engram_num_embeddings'][li]}") + token_map, cvocab = engram_compressed_token_map(self.dir_model if self.target_model_dir is None else self.target_model_dir) + if cvocab != hp["engram_compressed_vocab_size"]: + raise ValueError(f"compressed vocab {cvocab} != engram_compressed_vocab_size {hp['engram_compressed_vocab_size']}") + if len(token_map) != hp["vocab_size"]: + raise ValueError(f"token map has {len(token_map)} entries, vocab_size is {hp['vocab_size']}") + mults = engram_multipliers(layer_ids, max_ngram, cvocab) + w.add_engram_layers(layer_ids) + w.add_engram_num_embeddings(hp["engram_num_embeddings"]) + w.add_engram_max_ngram_size(max_ngram) + w.add_engram_vocab_size(hp["engram_vocab_size"]) + w.add_engram_head_count(n_heads) + w.add_engram_head_dim(hp["engram_head_dim"]) + w.add_engram_pad_token_id(hp["engram_pad_token_id"]) + w.add_engram_compressed_vocab_size(cvocab) + w.add_engram_token_map(token_map) + w.add_engram_multipliers([m for row in mults for m in row]) + w.add_engram_primes([p for layer in primes for per in layer for p in per]) + # per-layer row offsets of each (order, head) bucket range: exclusive cumsum of that layer's primes (order-major, + # head-minor) = NgramHashState.offsets; u64 so table addressing never narrows (tables are > 100 GB) + offsets = [] + for layer in primes: + flat = [p for per in layer for p in per] + offsets += [sum(flat[:i]) for i in range(len(flat))] + w.add_key_value(f"{self.gguf_writer.arch}.engram.offsets", offsets, gguf.GGUFValueType.ARRAY, gguf.GGUFValueType.UINT64) + logger.info(f"engram: compressed vocab {cvocab}, pad -> {token_map[hp['engram_pad_token_id']]}, multipliers {mults}") + + # DSpark (block layout only matters to the --mtp file, but the keys are cheap) + if "dspark_block_size" in hp: + w.add_dspark_block_size(hp["dspark_block_size"]) + w.add_dspark_noise_token_id(hp["dspark_noise_token_id"]) + w.add_dspark_target_layers(hp["dspark_target_layer_ids"]) + w.add_dspark_markov_rank(hp["dspark_markov_rank"]) + w.add_dspark_expert_count(hp["dspark_n_routed_experts"]) + w.add_dspark_expert_used_count(hp["dspark_num_experts_per_tok"]) + if self.mtp_only: + w.add_nextn_predict_layers(hp.get("num_nextn_predict_layers", 0)) + assert self.block_count == n_main + (hp.get("num_nextn_predict_layers", 0) if self.mtp_only else 0) + + # ---- weights ---------------------------------------------------------------------------------------------------- + + def dequant_model(self): + """FP8 (e4m3, ue8m0 scales) linears -> f32, exact. The block is derived per tensor from the scale shape and + checked against quantization_config.weight_block_size (V4 used 128, V4.1 uses 32x32).""" + fp8_dtypes = self._float8_dtypes() + bo_cfg, bi_cfg = self._fp8_block + to_remove: list[str] = [] + + def dequant(weight: Tensor, scale: Tensor, bo: int, bi: int) -> Tensor: + out_f, in_f = weight.shape + s = self._e8m0_to_float(scale) + s = s.repeat_interleave(bo, 0)[:out_f].repeat_interleave(bi, 1)[:, :in_f] + return weight.float() * s + + for name in list(self.model_tensors.keys()): + if not name.endswith(".scale") or self._ENGRAM_EMBED_RE.match(name): + continue + weight_name = name.removesuffix(".scale") + ".weight" + if weight_name not in self.model_tensors: + continue + wgen, sgen = self.model_tensors[weight_name], self.model_tensors[name] + w_meta = wgen() + if w_meta.dtype not in fp8_dtypes: + continue # MXFP4 experts (I8 packed) are handled in generate_extra_tensors + out_f, in_f = w_meta.shape + s_meta = sgen() + srows, scols = s_meta.shape + bo, bi = -(-out_f // srows), -(-in_f // scols) + if (bo, bi) != (bo_cfg, bi_cfg) or -(-out_f // bo) != srows or -(-in_f // bi) != scols: + raise ValueError(f"{weight_name}: FP8 block {(bo, bi)} from scale {tuple(s_meta.shape)} " + f"does not match weight_block_size {(bo_cfg, bi_cfg)}") + # exactness: e4m3 (|x| in [2^-9, 448]) * 2^(e-127) must stay inside BF16's range + _e8m0_check(LazyTorchTensor.to_eager(s_meta).view(torch.uint8).numpy(), name, 3, 246) + self.model_tensors[weight_name] = lambda w=wgen, s=sgen, bo=bo, bi=bi: dequant(w(), s(), bo, bi) + self._dsv4_fp8_dequantized.add(weight_name) + to_remove.append(name) + + for name in to_remove: + del self.model_tensors[name] + + def _write_mxfp4_experts(self, bid: int, n_experts: int) -> list[str]: + consumed: list[str] = [] + for proj, key in (("w1", gguf.MODEL_TENSOR.FFN_GATE_EXP), ("w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP), + ("w3", gguf.MODEL_TENSOR.FFN_UP_EXP)): + loaders, names = [], [] + for eid in range(n_experts): + wn, sn = f"layers.{bid}.ffn.experts.{eid}.{proj}.weight", f"layers.{bid}.ffn.experts.{eid}.{proj}.scale" + if wn not in self.model_tensors or sn not in self.model_tensors: + raise KeyError(f"Missing routed expert tensors {wn}") + loaders.append((self.model_tensors[wn], self.model_tensors[sn])) + names.append((wn, sn)) + consumed += [wn, sn] + extra = sorted(n for n in self.model_tensors + if re.match(rf"layers\.{bid}\.ffn\.experts\.(\d+)\.{proj}\.weight$", n) + and int(n.split(".")[4]) >= n_experts) + if extra: + raise KeyError(f"layer {bid}: unexpected experts beyond {n_experts}: {extra[:3]}") + data = self._mxfp4_expert_stream(names, loaders) + new_name = self.format_tensor_name(key, bid) + self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4) + return consumed + + @staticmethod + def _repack_mxfp4_np(packed: np.ndarray, scale: np.ndarray) -> np.ndarray: + """Same bytes as TextModel.repack_mxfp4_blocks (checked on the real experts), vectorised: + source byte k of a 32-group holds elements 2k (low) / 2k+1 (high); ggml byte j holds j (low) / j+16 (high).""" + rows, packed_cols = packed.shape + nb = packed_cols * 2 // 32 + if packed_cols * 2 % 32 != 0 or scale.shape != (rows, nb): + raise ValueError(f"MXFP4 shapes {packed.shape} / {scale.shape} do not match") + src = packed.reshape(rows, nb, 16) + lo_src, hi_src = src[:, :, :8], src[:, :, 8:] # elements 0..15 and 16..31 + out = np.empty((rows, nb, 17), dtype=np.uint8) + out[:, :, 0] = scale + q = out[:, :, 1:].reshape(rows, nb, 8, 2) + q[..., 0] = (lo_src & 0x0F) | ((hi_src & 0x0F) << 4) # j = 2k + q[..., 1] = (lo_src >> 4) | (hi_src & 0xF0) # j = 2k + 1 + return out.reshape(rows, nb * 17) + + def _mxfp4_expert_stream(self, names: list[tuple[str, str]], loaders) -> _StreamedBytes: + """Stacked [n_expert, rows, blocks*17] MXFP4 tensor, repacked by a thread pool at write time.""" + import threading + from concurrent.futures import ThreadPoolExecutor + from safetensors import safe_open + + rows, packed_cols = loaders[0][0]().shape + shape = (len(names), rows, packed_cols * 2 // 32 * 17) + n_main = self.hparams["num_hidden_layers"] + + def src_name(n: str) -> str: # undo the mtp.M -> layers.(n_main+M) rekey of filter_tensors + m = re.match(r"layers\.(\d+)\.(.+)$", n) + if m is not None and int(m.group(1)) >= n_main: + return f"mtp.{int(m.group(1)) - n_main}.{m.group(2)}" + return n + names = [(src_name(wn), src_name(sn)) for wn, sn in names] + parts = {n: self._tensor_part(n) for pair in names for n in pair} # weight and scale may sit in different shards + local = threading.local() + + def one(ws: tuple[str, str]) -> np.ndarray: + wn, sn = ws + hs = local.__dict__.setdefault("h", {}) + for part in (parts[wn], parts[sn]): + if part not in hs: + hs[part] = safe_open(part, framework="pt") + w = hs[parts[wn]].get_tensor(wn).view(torch.uint8).numpy() + sc = hs[parts[sn]].get_tensor(sn).view(torch.uint8).numpy() + return self._repack_mxfp4_np(w, sc) + + def chunks() -> Iterable[np.ndarray]: + with ThreadPoolExecutor(int(os.environ.get("DSV41_THREADS", "16"))) as ex: + yield from ex.map(one, names) + return _StreamedBytes(shape, chunks) + + def _engram_embed_tensor(self, bid: int) -> list[str]: + wn, sn = f"layers.{bid}.engram.embed.weight", f"layers.{bid}.engram.embed.scale" + wgen, sgen = self.model_tensors[wn], self.model_tensors[sn] + rows, dim = wgen().shape + if tuple(sgen().shape) != (rows, dim // 32) or dim % 32 != 0: + raise ValueError(f"{wn}: scale shape {tuple(sgen().shape)} is not 1x32 blocks of {(rows, dim)}") + li = list(self.hparams["engram_layer_ids"]).index(bid) + if rows != self.hparams["engram_num_embeddings"][li]: + raise ValueError(f"{wn}: {rows} rows != engram_num_embeddings {self.hparams['engram_num_embeddings'][li]}") + new_name = self.format_tensor_name(gguf.MODEL_TENSOR.ENGRAM_EMBD, bid) + wpart, spart = self._tensor_part(wn), self._tensor_part(sn) + chunk_rows = 1 << 20 + + def raw_chunks() -> Iterable[tuple[np.ndarray, np.ndarray]]: + from safetensors import safe_open + with safe_open(wpart, framework="pt") as fw, safe_open(spart, framework="pt") as fs: + ws, ss = fw.get_slice(wn), fs.get_slice(sn) + for r0 in range(0, rows, chunk_rows): + r1 = min(rows, r0 + chunk_rows) + q = ws[r0:r1].view(torch.uint8).numpy().reshape(-1, 32) + e = ss[r0:r1].view(torch.uint8).numpy().reshape(-1, 1) + yield q, e + + if self._engram_type == "mxfp8": + def chunks() -> Iterable[np.ndarray]: + for q, e in raw_chunks(): + if e.max() == 0xFF: + raise ValueError(f"{wn}: E8M0 NaN scale (0xFF)") + if ((q & 0x7F) == 0x7F).any(): + raise ValueError(f"{wn}: E4M3 NaN code") + yield np.concatenate([e, q], axis=1) + data = _StreamedBytes((rows, dim // 32 * 33), chunks) + self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP8) # type: ignore[arg-type] + else: + to_bf16 = self._engram_type == "bf16" + + def chunks() -> Iterable[np.ndarray]: + for q, e in raw_chunks(): + _e8m0_check(e, wn, 3, 246) + v = torch.from_numpy(q).view(torch.float8_e4m3fn).float() * torch.exp2(torch.from_numpy(e).float() - 127) + v = v.reshape(-1, dim) + yield (v.bfloat16().view(torch.uint8) if to_bf16 else v.view(torch.uint8)).numpy() + esz = 2 if to_bf16 else 4 + data = _StreamedBytes((rows, dim * esz), chunks) + qt = gguf.GGMLQuantizationType.BF16 if to_bf16 else gguf.GGMLQuantizationType.F32 + self.gguf_writer.add_tensor(new_name, data, raw_shape=(rows, dim), raw_dtype=qt) # type: ignore[arg-type] + logger.info(f"{new_name}: Engram table {rows} x {dim} copied as {self._engram_type.upper()}") + return [wn, sn] + + _weight_map: dict[str, str] | None = None + + def _tensor_part(self, name: str) -> Path: + if self._weight_map is None: + import json + with open(self.dir_model / "model.safetensors.index.json", "r", encoding="utf-8") as f: + self._weight_map = json.load(f)["weight_map"] + return self.dir_model / self._weight_map[name] + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + if self._dsv4_mxfp4_generated: + return () + hp = self.hparams + consumed: list[str] = [] + n_main = hp["num_hidden_layers"] + for bid in range(self.block_count): + if self.mtp_only and bid < n_main: + continue + n_exp = hp["n_routed_experts"] if bid < n_main else hp["dspark_n_routed_experts"] + consumed += self._write_mxfp4_experts(bid, n_exp) + if not self.mtp_only: + for bid in hp["engram_layer_ids"]: + consumed += self._engram_embed_tensor(bid) + for name in consumed: + del self.model_tensors[name] + self._dsv4_mxfp4_generated = True + return () + + _LAYER_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = { + "attn.indexer.wk.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_K, ".weight"), + "attn.indexer.k_norm.weight": (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".weight"), + "engram.wkv.weight": (gguf.MODEL_TENSOR.ENGRAM_KV, ".weight"), + "main_proj.weight": (gguf.MODEL_TENSOR.DSPARK_MAIN_PROJ, ".weight"), + "main_norm.weight": (gguf.MODEL_TENSOR.DSPARK_MAIN_NORM, ".weight"), + "norm.weight": (gguf.MODEL_TENSOR.DSPARK_NORM, ".weight"), + } + _ROOT_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = { + "dspark.markov_head.embed.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W1, ".weight"), + "dspark.markov_head.head.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W2, ".weight"), + "dspark.confidence_head.proj.weight": (gguf.MODEL_TENSOR.DSPARK_CONF_PROJ, ".weight"), + } + + def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]: + if name in self._ROOT_MAP: + return self._ROOT_MAP[name] + m = re.match(r"layers\.(\d+)\.(.+)$", name) + if m is not None and m.group(2) in self._LAYER_MAP: + return self._LAYER_MAP[m.group(2)] + key, suffix = super()._map_dsv4_tensor_name(name, bid) + if key in (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_APE, gguf.MODEL_TENSOR.FFN_GATE_TID2EID) or \ + key.name.startswith(("INDEXER_COMPRESSOR_", "NEXTN_", "HC_HEAD_")): + raise ValueError(f"{name!r} is a DeepSeek-V4 tensor that V4.1 does not have") + return key, suffix + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if re.match(r"layers\.\d+\.ffn\.experts\.\d+\.w[123]\.(weight|scale)$", name): + return [] + if re.match(r"layers\.\d+\.engram\.k_weight$", name): + return [] # folded into engram_gate together with q_weight + if (m := re.match(r"layers\.(\d+)\.engram\.q_weight$", name)) is not None: + k = self.model_tensors.get(f"layers.{m.group(1)}.engram.k_weight") + if k is None: + raise KeyError(f"missing layers.{m.group(1)}.engram.k_weight") + # bf16 * bf16 is exact in f32 + gate = data_torch.float() * k().float() + return [(self.format_tensor_name(gguf.MODEL_TENSOR.ENGRAM_GATE, bid), gate)] + key, suffix = self._map_dsv4_tensor_name(name, bid) + new_name = self.format_tensor_name(key, None if name.startswith("dspark.") else bid, suffix) + return [(new_name, data_torch)] + + def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: + # router, engram gate and confidence head stay F32 (exact from bf16; verify_types.py pins the router to F32) + if new_name.endswith((".engram_gate.weight", ".ffn_gate_inp.weight")) or new_name == "conf_proj.weight": + return gguf.GGMLQuantizationType.F32 + dense_qt = {"q8_0": gguf.GGMLQuantizationType.Q8_0, "f32": gguf.GGMLQuantizationType.F32}.get( + self._dense_type, gguf.GGMLQuantizationType.BF16) + if name in self._dsv4_fp8_dequantized and n_dims >= 2: + return dense_qt + if name in self._dsv4_f32_tensors: + return gguf.GGMLQuantizationType.F32 + if name in self._dsv4_bf16_tensors and n_dims >= 2: + return dense_qt + return False + + + +class DeepseekV41DSparkModel(DeepseekV41Model): + """DeepSeek-V4.1 DSpark speculative draft (mtp.0..2 of the checkpoint) as a DFLASH GGUF (`--dspark`). + + Layout = the V4 DSpark draft (conversion/deepseek.py DeepseekV4DSparkModel) with V4.1's differences flagged by + `dflash.hyper_connection.single_pass`: stages blk.0..2 (128 experts top-3, MXFP4 bit-exact, dense BF16), + mtp.0.main_proj / main_norm -> fc / enc.output_norm, mtp.2.norm -> output_norm, markov_head.embed / head -> + markov_w1 / markov_w2, confidence_head.proj -> conf_proj (F32), no hc_head. Token embeddings and the lm head are + the target's (shared at runtime through ctx_other). target_layers are the V4.1 ids as-is: the draft reads the + attention input of those layers (model.py Transformer.forward), which is what llama_set_embeddings_layer_inp + extracts. + """ + model_arch = gguf.MODEL_ARCH.DFLASH + mtp_only = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + hp = self.hparams + self._n_main = hp["num_hidden_layers"] + self._n_stage = hp["num_nextn_predict_layers"] + if self._n_stage <= 0 or "dspark_block_size" not in hp: + raise ValueError("DeepSeek-V4.1 DSpark export needs num_nextn_predict_layers > 0 and the dspark_* config") + # the main-model rekey (mtp.M -> layers.n_main+M) is kept for tensor selection and the MXFP4 expert source + # names; format_tensor_name() maps the block ids back to 0..n_stage-1 + self.block_count = self._n_main + self._n_stage + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self._n_stage) + + def format_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None = None, suffix: str = ".weight") -> str: + if bid is not None and bid >= self._n_main: + bid -= self._n_main + return super().format_tensor_name(key, bid, suffix) + + def set_vocab(self): + super().set_vocab() + self.gguf_writer.add_mask_token_id(self.hparams["dspark_noise_token_id"]) + + def set_gguf_parameters(self): + hp = self.hparams + n = self._n_stage + w = self.gguf_writer + w.add_block_count(n) + w.add_context_length(hp["max_position_embeddings"]) + w.add_embedding_length(hp["hidden_size"]) + w.add_feed_forward_length(hp["moe_intermediate_size"]) + w.add_head_count(hp["num_attention_heads"]) + w.add_head_count_kv(1) + w.add_key_length(hp["head_dim"]) + w.add_value_length(hp["head_dim"]) + w.add_rope_dimension_count(hp["qk_rope_head_dim"]) + # DSpark stages are compress-ratio-0 layers: plain rope at rope_theta, no YaRN + w.add_rope_freq_base(float(hp["rope_theta"])) + w.add_layer_norm_rms_eps(hp["rms_norm_eps"]) + w.add_q_lora_rank(hp["q_lora_rank"]) + w.add_sliding_window(hp["sliding_window"]) + w.add_expert_count(hp["dspark_n_routed_experts"]) + w.add_expert_used_count(hp["dspark_num_experts_per_tok"]) + w.add_expert_feed_forward_length(hp["moe_intermediate_size"]) + w.add_expert_shared_count(hp["n_shared_experts"]) + w.add_expert_weights_scale(hp["routed_scaling_factor"]) + w.add_expert_weights_norm(hp["norm_topk_prob"]) + w.add_expert_gating_func(gguf.ExpertGatingFuncType.SQRTSOFTPLUS) + w.add_swiglu_clamp_exp([hp["swiglu_limit"]] * n) + w.add_swiglu_clamp_shexp([hp["swiglu_limit"]] * n) + w.add_attention_output_group_count(hp["o_groups"]) + w.add_attention_output_lora_rank(hp["o_lora_rank"]) + w.add_attention_compress_ratios([0] * n) + w.add_hyper_connection_count(hp["hc_mult"]) + w.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"]) + w.add_hyper_connection_epsilon(hp["hc_eps"]) + w.add_hyper_connection_single_pass(True) + w.add_block_size(hp["dspark_block_size"]) + w.add_target_layers(hp["dspark_target_layer_ids"]) + w.add_sample_from_anchor(True) + w.add_has_confidence_head(True) + w.add_file_type(self.ftype) + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + # stage experts only (no Engram tables in the draft) + return super().generate_extra_tensors() + + _DSPARK_ROOT: dict[str, gguf.MODEL_TENSOR] = { + "main_proj.weight": gguf.MODEL_TENSOR.FC, + "main_norm.weight": gguf.MODEL_TENSOR.ENC_OUTPUT_NORM, + "norm.weight": gguf.MODEL_TENSOR.OUTPUT_NORM, + } + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name in ("embed.weight", "head.weight", "norm.weight"): + return [] # the target's (ctx_other) + if name.endswith(".ffn.gate.bias_vl"): + return [] # drafts are text-only + m = re.match(r"layers\.(\d+)\.(main_proj\.weight|main_norm\.weight|norm\.weight)$", name) + if m is not None and int(m.group(1)) >= self._n_main: + return [(self.format_tensor_name(self._DSPARK_ROOT[m.group(2)], None), data_torch)] + return super().modify_tensors(data_torch, name, bid) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index e09616b190cf..5e263fd16fa7 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -267,11 +267,15 @@ def main() -> None: sys.exit(1) if args.dspark: - if is_mistral_format or model_architecture != "DeepseekV4ForCausalLM": - logger.error("--dspark is only supported for DeepseekV4ForCausalLM") + if is_mistral_format or model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"): + logger.error("--dspark is only supported for DeepseekV4ForCausalLM and DeepseekV41ForCausalLM") sys.exit(1) - from conversion.deepseek import DeepseekV4DSparkModel - model_class = DeepseekV4DSparkModel + if model_architecture == "DeepseekV41ForCausalLM": + from conversion.deepseek41 import DeepseekV41DSparkModel + model_class = DeepseekV41DSparkModel + else: + from conversion.deepseek import DeepseekV4DSparkModel + model_class = DeepseekV4DSparkModel if args.mtp or args.no_mtp: if not model_class.supports_mtp_export: diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 77e0681b64a9..5027d5a83cd2 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -436,7 +436,8 @@ extern "C" { GGML_TYPE_IQ1_XS = 64, // 1.4375 bpw, 1024-entry grid GGML_TYPE_IQ1_XXS = 65, // 1.3125 bpw, 512-entry grid GGML_TYPE_IQ1_XXXS = 66, // 1.1875 bpw, 256-entry grid - GGML_TYPE_COUNT = 67, + GGML_TYPE_MXFP8 = 67, // fork-local, MXFP8 E4M3 (32 elements, E8M0 scale) + GGML_TYPE_COUNT = 68, }; // [TAG_GGML_PREC] diff --git a/ggml/src/ggml-common.h b/ggml/src/ggml-common.h index ddb5355e88d8..bf214e2f9d8a 100644 --- a/ggml/src/ggml-common.h +++ b/ggml/src/ggml-common.h @@ -225,6 +225,14 @@ typedef struct { } block_mxfp4; static_assert(sizeof(block_mxfp4) == sizeof(uint8_t) + QK_MXFP4/2, "wrong mxfp4 block size/padding"); +// MXFP8 (OCP MX): 32 E4M3 (fn) values sharing one E8M0 scale, value = e4m3(q[j]) * 2^(e - 127) +#define QK_MXFP8 32 +typedef struct { + uint8_t e; // E8M0 + uint8_t qs[QK_MXFP8]; // E4M3 (fn) codes +} block_mxfp8; +static_assert(sizeof(block_mxfp8) == sizeof(uint8_t) + QK_MXFP8, "wrong mxfp8 block size/padding"); + #define QK_NVFP4 64 #define QK_NVFP4_SUB 16 // sub-block size for per-group scales typedef struct { diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 92649a0ebca3..a571fc2aec03 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -429,6 +429,16 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return true; } + // MXFP8 is a gather-only type: GET_ROWS reads it, nothing else does (and nothing writes it) + if (op->type == GGML_TYPE_MXFP8) { + return false; + } + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (op->src[i] && op->src[i]->type == GGML_TYPE_MXFP8 && !(op->op == GGML_OP_GET_ROWS && i == 0)) { + return false; + } + } + // check extra buffer types // note: only the first sources are checked for extra buffer types to reduce overhead, increase if necessary for (int i = 0; i < 4; i++) { diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index aedd6f6d71fd..132975f2c68f 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -5471,6 +5471,7 @@ void ggml_compute_forward_get_rows( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_MXFP8: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: diff --git a/ggml/src/ggml-cuda/dequantize.cuh b/ggml/src/ggml-cuda/dequantize.cuh index 651524be7b3e..dc2cbe3d75ec 100644 --- a/ggml/src/ggml-cuda/dequantize.cuh +++ b/ggml/src/ggml-cuda/dequantize.cuh @@ -119,6 +119,24 @@ static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const in v.y *= d; } +// E4M3 (fn) decode, exact in fp32; 0x7F/0xFF are NaN +static __device__ __forceinline__ float ggml_cuda_e4m3_to_fp32(const uint8_t x) { + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + float v = exp == 0 ? ldexpf((float) man, -9) : ldexpf((float) (8 + man), exp - 10); + v = (exp == 0xF && man == 0x7) ? __int_as_float(0x7FC00000) : v; + return (x & 0x80) ? -v : v; +} + +static __device__ __forceinline__ void dequantize_mxfp8(const void * vx, const int64_t ib, const int iqs, float2 & v){ + const block_mxfp8 * x = (const block_mxfp8 *) vx; + + const float d = x[ib].e == 0xFF ? __int_as_float(0x7FC00000) : ggml_cuda_e8m0_to_fp32(x[ib].e); + + v.x = ggml_cuda_e4m3_to_fp32(x[ib].qs[iqs + 0]) * d; + v.y = ggml_cuda_e4m3_to_fp32(x[ib].qs[iqs + 1]) * d; +} + //================================== k-quants // Each call dequantizes one super-block of QK_K values into y using the diff --git a/ggml/src/ggml-cuda/getrows.cu b/ggml/src/ggml-cuda/getrows.cu index d97e8fdb0f40..95f6d2383a78 100644 --- a/ggml/src/ggml-cuda/getrows.cu +++ b/ggml/src/ggml-cuda/getrows.cu @@ -344,6 +344,10 @@ static void ggml_cuda_get_rows_switch_src0_type( get_rows_cuda_q(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); break; + case GGML_TYPE_MXFP8: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; case GGML_TYPE_Q2_K: get_rows_cuda_kq<64, dst_t, dequantize_q2_K>(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 9b6366addcfe..9b0ff9397a77 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1965,6 +1965,25 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor // Therefore, in such cases use cuBLAS. const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; + + // [TAG_GGML_PREC] src1 requested at GGML_PREC_F32 with F32 weights: keep both operands in F32, i.e. no TF32 + // rounding (the MMF tensor-core path and cuBLAS tensor-op math both round to a 10-bit mantissa). Small batches use + // the F32 matrix-vector kernel, larger ones cuBLAS SGEMM with the default (non-TF32) math mode. + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 && + ggml_get_op_params_i32(dst, 3) == GGML_PREC_F32) { + const int cc_f32 = ggml_cuda_info().devices[ctx.device].cc; + const int warp_size_f32 = ggml_cuda_info().devices[ctx.device].warp_size; + if (!bad_padding_clear && ne11 <= MMVF_MAX_BATCH_SIZE && + ggml_cuda_should_use_mmvf(src0->type, cc_f32, warp_size_f32, src0->ne, src0->nb, /*ne11 =*/ 1)) { + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + return; + } + CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(), CUBLAS_DEFAULT_MATH)); + ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(), CUBLAS_TF32_TENSOR_OP_MATH)); + return; + } + if (bad_padding_clear || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); return; @@ -5648,6 +5667,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_IQ1_S: case GGML_TYPE_IQ1_M: case GGML_TYPE_IQ4_XS: + case GGML_TYPE_MXFP8: return true; case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index b0f60999021f..08a48132265c 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -92,14 +92,12 @@ struct top_k_radix_state { uint32_t prefix; uint32_t prefix_mask; int rank; - int greater_count; - int equal_count; }; static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) { const int row = blockIdx.x * blockDim.x + threadIdx.x; if (row < nrows) { - states[row] = {0, 0, k, 0, 0}; + states[row] = {0, 0, k}; } } @@ -170,11 +168,77 @@ static __global__ void top_k_radix_select( } } -static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) { - const int row = blockIdx.x * blockDim.x + threadIdx.x; - if (row < nrows) { - states[row].greater_count = 0; - states[row].equal_count = 0; +// Block-wide inclusive prefix count of `flag` (0/1) over the threads of a block; returns this thread's inclusive count +// and the block total. Must be called by every thread of the block. +template +static __device__ __forceinline__ int top_k_block_prefix_count(const int flag, int & total) { + static_assert(BLOCK_SIZE % WARP_SIZE == 0, "BLOCK_SIZE must be a multiple of WARP_SIZE"); + constexpr int NWARPS = BLOCK_SIZE / WARP_SIZE; + __shared__ int warp_sums[NWARPS]; + + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; + + int incl = flag; +#pragma unroll + for (int off = 1; off < WARP_SIZE; off *= 2) { + const int y = __shfl_up_sync(0xFFFFFFFF, incl, off, WARP_SIZE); + if (lane >= off) { + incl += y; + } + } + if (lane == WARP_SIZE - 1) { + warp_sums[warp] = incl; + } + __syncthreads(); + int warp_off = 0; + total = 0; + for (int w = 0; w < NWARPS; ++w) { + warp_off += w < warp ? warp_sums[w] : 0; + total += warp_sums[w]; + } + __syncthreads(); + return warp_off + incl; +} + +// Elements above the cut are written in ascending column order (not atomicAdd order), so the index order of the +// result is reproducible run to run. Consumers that gather rows in top-k order (sparse attention over the selected +// rows) otherwise accumulate in a different order every run. Each block owns a contiguous column chunk: pass one +// counts its elements above the cut, pass two writes them at the sum of the earlier chunks' counts plus an ordered +// in-block prefix count. +template +static __global__ void top_k_radix_count_greater( + const float * __restrict__ src, + const top_k_radix_state * __restrict__ states, + int * __restrict__ chunk_counts, + int ncols, + int chunk_size, + int blocks_per_row) { + const int row = blockIdx.x / blocks_per_row; + const int row_block = blockIdx.x % blocks_per_row; + const float * row_src = src + (size_t) row * ncols; + const uint32_t prefix = states[row].prefix; + + const int64_t col_begin = (int64_t) row_block * chunk_size; + const int64_t col_end = col_begin + chunk_size < ncols ? col_begin + chunk_size : (int64_t) ncols; + + int count = 0; + for (int64_t col = col_begin + threadIdx.x; col < col_end; col += BLOCK_SIZE) { + count += top_k_float_to_ordered(row_src[col]) > prefix; + } + count = warp_reduce_sum(count); + + __shared__ int warp_counts[BLOCK_SIZE / WARP_SIZE]; + if (threadIdx.x % WARP_SIZE == 0) { + warp_counts[threadIdx.x / WARP_SIZE] = count; + } + __syncthreads(); + if (threadIdx.x == 0) { + int total = 0; + for (int w = 0; w < BLOCK_SIZE / WARP_SIZE; ++w) { + total += warp_counts[w]; + } + chunk_counts[blockIdx.x] = total; } } @@ -182,30 +246,64 @@ template static __global__ void top_k_radix_gather( const float * __restrict__ src, int * __restrict__ dst, - top_k_radix_state * __restrict__ states, + const top_k_radix_state * __restrict__ states, + const int * __restrict__ chunk_counts, int ncols, int k, + int chunk_size, int blocks_per_row) { const int row = blockIdx.x / blocks_per_row; const int row_block = blockIdx.x % blocks_per_row; - const int tid = threadIdx.x; const float * row_src = src + (size_t) row * ncols; int * row_dst = dst + (size_t) row * k; - top_k_radix_state * state = &states[row]; + const uint32_t prefix = states[row].prefix; - for (int64_t col = row_block * BLOCK_SIZE + tid; - col < ncols; - col += blocks_per_row * BLOCK_SIZE) { - const uint32_t key = top_k_float_to_ordered(row_src[col]); - if (key > state->prefix) { - const int pos = atomicAdd(&state->greater_count, 1); - row_dst[pos] = col; - } else if (key == state->prefix) { - const int pos = atomicAdd(&state->equal_count, 1); - if (pos < state->rank) { - row_dst[k - state->rank + pos] = col; - } + int base = 0; + for (int j = 0; j < row_block; ++j) { + base += chunk_counts[row * blocks_per_row + j]; + } + + const int64_t col_begin = (int64_t) row_block * chunk_size; + const int64_t col_end = col_begin + chunk_size < ncols ? col_begin + chunk_size : (int64_t) ncols; + for (int64_t col0 = col_begin; col0 < col_end; col0 += BLOCK_SIZE) { + const int64_t col = col0 + threadIdx.x; + const int gt = col < col_end && top_k_float_to_ordered(row_src[col]) > prefix; + int total; + const int incl = top_k_block_prefix_count(gt, total); + if (gt) { + row_dst[base + incl - 1] = col; + } + base += total; + } +} + +// Elements equal to the k-th key fill the last `rank` slots, taking the lowest column indices first. With the atomic +// slot order used for the greater elements, which of several tied elements made the cut changed from run to run +// (exactly tied scores are common, e.g. masked or ReLU-zeroed rows), so the selected set was not reproducible. +// One block per row scans the columns in order with a block-wide prefix count, stopping once the slots are full. +template +static __global__ void top_k_radix_gather_equal( + const float * __restrict__ src, + int * __restrict__ dst, + const top_k_radix_state * __restrict__ states, + int ncols, + int k) { + const int row = blockIdx.x; + const int tid = threadIdx.x; + const float * row_src = src + (size_t) row * ncols; + int * row_dst = dst + (size_t) row * k; + const top_k_radix_state state = states[row]; + + int base = 0; + for (int64_t col0 = 0; col0 < ncols && base < state.rank; col0 += BLOCK_SIZE) { + const int64_t col = col0 + tid; + const int eq = col < ncols && top_k_float_to_ordered(row_src[col]) == state.prefix; + int total; + const int pos = base + top_k_block_prefix_count(eq, total) - eq; + if (eq && pos < state.rank) { + row_dst[k - state.rank + pos] = col; } + base += total; } } @@ -216,6 +314,7 @@ static void top_k_radix_cuda( constexpr int RADIX_BITS = 8; constexpr int NBINS = 1 << RADIX_BITS; const int blocks_per_row = (int) std::min(((int64_t) ncols + 1023) / 1024, 64); + const int chunk_size = (int) (((int64_t) ncols + blocks_per_row - 1) / blocks_per_row); // chunk the rows to bound the histogram memory to 64 MB const int64_t chunk_nrows = ggml_cuda_chunk_nrows((size_t) blocks_per_row * NBINS * sizeof(int), nrows); @@ -239,11 +338,14 @@ static void top_k_radix_cuda( <<>>(histograms, states, blocks_per_row, shift); } - top_k_radix_reset_counters - <<<(iter_nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, iter_nrows); + // the last select pass has consumed the histograms: reuse them for the per-chunk counts + top_k_radix_count_greater + <<>>(src, states, histograms, ncols, chunk_size, blocks_per_row); top_k_radix_gather <<>>( - src, dst, states, ncols, k, blocks_per_row); + src, dst, states, histograms, ncols, k, chunk_size, blocks_per_row); + top_k_radix_gather_equal + <<>>(src, dst, states, ncols, k); src += (size_t) ncols * iter_nrows; dst += (size_t) k * iter_nrows; diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index ae26e0c23b46..96ab08ae49f0 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -564,6 +564,51 @@ static inline uint8_t ggml_fp32_to_ue4m3(float x) { return (uint8_t) ((ue4m3_exp << 3) | ue4m3_man); } +// E4M3 (fn): signed, 4 exp bits (bias=7), 3 mantissa bits, no inf, 0x7F/0xFF = NaN (used by MXFP8) +// every E4M3 value is exactly representable in fp32 +static inline float ggml_e4m3_to_fp32(uint8_t x) { + const int sign = x & 0x80; + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + float v; + if (exp == 0xF && man == 0x7) { + v = NAN; + } else if (exp == 0) { + v = ldexpf((float) man, -9); + } else { + v = ldexpf((float) (8 + man), exp - 10); + } + return sign ? -v : v; +} + +// fp32 -> E4M3 (fn), round to nearest even, saturating to +-448, NaN -> 0x7F +static inline uint8_t ggml_fp32_to_e4m3(float x) { + if (x != x) { + return 0x7F; + } + const uint8_t sign = signbit(x) ? 0x80 : 0x00; + float a = fabsf(x); + if (a >= 448.0f) { + return sign | 0x7E; + } + if (a < 0.015625f) { // < 2^-6: subnormal grid 2^-9 (m == 8 is the smallest normal, code 0x08) + const int m = (int) nearbyintf(a * 512.0f); + return sign | (uint8_t) m; + } + int e; + const float f = frexpf(a, &e); // a = f * 2^e, f in [0.5, 1) + int m = (int) nearbyintf((f * 2.0f - 1.0f) * 8.0f); + int E = e - 1 + 7; + if (m == 8) { + m = 0; + E++; + } + if (E > 15 || (E == 15 && m == 7)) { + return sign | 0x7E; + } + return sign | (uint8_t) ((E << 3) | m); +} + /** * Converts brain16 to float32. * diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c index 776b2effb63d..dc9469156162 100644 --- a/ggml/src/ggml-quants.c +++ b/ggml/src/ggml-quants.c @@ -347,6 +347,35 @@ static inline int best_index_mxfp4(float x, float e) { return best_index; } +void quantize_row_mxfp8_ref(const float * GGML_RESTRICT x, block_mxfp8 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_MXFP8; + + assert(k % qk == 0); + + const int64_t nb = k / qk; + + for (int64_t i = 0; i < nb; i++) { + float amax = 0.0f; + for (int j = 0; j < qk; j++) { + amax = MAX(amax, fabsf(x[i*qk + j])); + } + + // smallest power of two scale with amax/scale <= 448 (largest e4m3 value) + int e = 0; + if (amax > 0.0f) { + e = (int) ceilf(log2f(amax / 448.0f)) + 127; + e = MAX(0, MIN(254, e)); + } + const float d = GGML_E8M0_TO_FP32((uint8_t) e); + const float id = 1.0f/d; + + y[i].e = (uint8_t) e; + for (int j = 0; j < qk; ++j) { + y[i].qs[j] = ggml_fp32_to_e4m3(x[i*qk + j]*id); + } + } +} + void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k) { static const int qk = QK_MXFP4; @@ -586,6 +615,23 @@ void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_REST } } +void dequantize_row_mxfp8(const block_mxfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK_MXFP8; + + assert(k % qk == 0); + + const int64_t nb = k / qk; + + for (int64_t i = 0; i < nb; i++) { + // exact: e4m3 * 2^(e-127) is representable in fp32 for every code and scale + const float d = x[i].e == 0xFF ? NAN : GGML_E8M0_TO_FP32(x[i].e); + + for (int j = 0; j < qk; ++j) { + y[i*qk + j] = ggml_e4m3_to_fp32(x[i].qs[j])*d; + } + } +} + void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { static const int qk = QK_NVFP4; static const int qk_sub = QK_NVFP4_SUB; @@ -2306,6 +2352,12 @@ size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, return nrow * ggml_row_size(GGML_TYPE_MXFP4, n_per_row); } +size_t quantize_mxfp8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + GGML_UNUSED(quant_weights); + quantize_row_mxfp8_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_MXFP8, n_per_row); +} + size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { GGML_UNUSED(quant_weights); quantize_row_nvfp4_ref(src, dst, (int64_t)nrow*n_per_row); @@ -5890,6 +5942,10 @@ bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbyte { VALIDATE_ROW_DATA_E_E8M0_IMPL(block_mxfp4, data, nb); } break; + case GGML_TYPE_MXFP8: + { + VALIDATE_ROW_DATA_E_E8M0_IMPL(block_mxfp8, data, nb); + } break; case GGML_TYPE_NVFP4: { // UE4M3 scales are uint8_t — all byte values are valid diff --git a/ggml/src/ggml-quants.h b/ggml/src/ggml-quants.h index 369ad5da5307..d76fcefc2b4b 100644 --- a/ggml/src/ggml-quants.h +++ b/ggml/src/ggml-quants.h @@ -24,6 +24,7 @@ GGML_API void quantize_row_q8_0_ref(const float * GGML_RESTRICT x, block_q8_0 * GGML_API void quantize_row_q8_1_ref(const float * GGML_RESTRICT x, block_q8_1 * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_mxfp8_ref(const float * GGML_RESTRICT x, block_mxfp8 * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_nvfp4_ref(const float * GGML_RESTRICT x, block_nvfp4 * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_q2_K_ref(const float * GGML_RESTRICT x, block_q2_K * GGML_RESTRICT y, int64_t k); @@ -53,6 +54,7 @@ GGML_API void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GG //GGML_API void dequantize_row_q8_1(const block_q8_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_mxfp8(const block_mxfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); @@ -109,6 +111,7 @@ GGML_API size_t quantize_q5_1(const float * GGML_RESTRICT src, void * GGML_RESTR GGML_API size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_mxfp8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API void iq2xs_init_impl(enum ggml_type type); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 199af8e21cdf..f83f575fb349 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -768,6 +768,14 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .to_float = (ggml_to_float_t) dequantize_row_mxfp4, .from_float_ref = (ggml_from_float_t)quantize_row_mxfp4_ref, }, + [GGML_TYPE_MXFP8] = { + .type_name = "mxfp8", + .blck_size = QK_MXFP8, + .type_size = sizeof(block_mxfp8), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_mxfp8, + .from_float_ref = (ggml_from_float_t)quantize_row_mxfp8_ref, + }, [GGML_TYPE_NVFP4] = { .type_name = "nvfp4", .blck_size = QK_NVFP4, @@ -8146,6 +8154,7 @@ size_t ggml_quantize_chunk( case GGML_TYPE_Q5_1: result = quantize_q5_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q8_0: result = quantize_q8_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_MXFP4: result = quantize_mxfp4 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_MXFP8: result = quantize_mxfp8 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_NVFP4: result = quantize_nvfp4 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q2_K: result = quantize_q2_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q3_K: result = quantize_q3_K (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 7c8efe7dfa3d..7504f81d5e31 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -211,6 +211,7 @@ class Attention: VALUE_SCALE = "{arch}.attention.value_scale" COMPRESS_RATIOS = "{arch}.attention.compress_ratios" COMPRESS_ROPE_FREQ_BASE = "{arch}.attention.compress_rope_freq_base" + KV_SOURCE_LAYERS = "{arch}.attention.kv_source_layers" # deepseek41 TEMPERATURE_LENGTH = "{arch}.attention.temperature_length" KEY_LENGTH_MLA = "{arch}.attention.key_length_mla" VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla" @@ -236,6 +237,10 @@ class Indexer: TYPES = "{arch}.attention.indexer.types" KPOOL = "{arch}.attention.indexer.kpool" # GLM5-Next KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail" # GLM5-Next + SOURCE_LAYERS = "{arch}.attention.indexer.source_layers" # deepseek41 + CANDIDATE_SOURCE_LAYER = "{arch}.attention.indexer.candidate_source_layer" # deepseek41, -1 = off + CANDIDATE_TOP_K_BLOCKS = "{arch}.attention.indexer.candidate_top_k_blocks" # deepseek41 + CANDIDATE_BLOCK_SIZE = "{arch}.attention.indexer.candidate_block_size" # deepseek41 class HyperConnection: COUNT = "{arch}.hyper_connection.count" @@ -245,6 +250,29 @@ class HyperConnection: MAGNITUDE = "{arch}.hyper_connection.magnitude" # absent means the mix projection is full rank (DeepSeek-V4 behaviour) LOW_RANK = "{arch}.hyper_connection.low_rank" + # single-pass mHC: each sublayer collapses with the previous sublayer's pre-mix (DeepSeek-V4.1) + SINGLE_PASS = "{arch}.hyper_connection.single_pass" + + class Engram: # deepseek41 hashed n-gram memory + LAYERS = "{arch}.engram.layers" # u32[n_engram] + NUM_EMBEDDINGS = "{arch}.engram.num_embeddings" # u32[n_engram], table rows per layer + MAX_NGRAM_SIZE = "{arch}.engram.max_ngram_size" # u32, n-gram orders 2..max + VOCAB_SIZE = "{arch}.engram.vocab_size" # u32, prime search start (informational) + HEAD_COUNT = "{arch}.engram.head_count" # u32, heads per n-gram order + HEAD_DIM = "{arch}.engram.head_dim" # u32, row width + PAD_TOKEN_ID = "{arch}.engram.pad_token_id" # u32, raw token id filling missing history + COMPRESSED_VOCAB_SIZE = "{arch}.engram.compressed_vocab_size" # u32 + TOKEN_MAP = "{arch}.engram.token_map" # i32[n_vocab], raw id -> compressed id + MULTIPLIERS = "{arch}.engram.multipliers" # u64[n_engram * max_ngram], layer-major + PRIMES = "{arch}.engram.primes" # u32[n_engram * (max_ngram-1) * head_count], layer, order, head + + class DSpark: # deepseek41 draft head (mtp.* blocks) + BLOCK_SIZE = "{arch}.dspark.block_size" + NOISE_TOKEN_ID = "{arch}.dspark.noise_token_id" + TARGET_LAYERS = "{arch}.dspark.target_layers" + MARKOV_RANK = "{arch}.dspark.markov_rank" + EXPERT_COUNT = "{arch}.dspark.expert_count" + EXPERT_USED_COUNT = "{arch}.dspark.expert_used_count" class PerLayerEmbedding: LAYERS = "{arch}.ple.layers" @@ -626,6 +654,7 @@ class MODEL_ARCH(IntEnum): DEEPSEEK2OCR = auto() DEEPSEEK32 = auto() DEEPSEEK4 = auto() + DEEPSEEK41 = auto() CHATGLM = auto() GLM4 = auto() GLM4_MOE = auto() @@ -1302,6 +1331,13 @@ class MODEL_TENSOR(IntEnum): DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed DSPARK_MARKOV_W2 = auto() # markov head: bias projection DSPARK_CONF_PROJ = auto() # confidence head + DSPARK_MAIN_PROJ = auto() # deepseek41: target hidden taps -> n_embd + DSPARK_MAIN_NORM = auto() # deepseek41 + DSPARK_NORM = auto() # deepseek41: final norm of the last DSpark stage + # engram (deepseek41) + ENGRAM_EMBD = auto() + ENGRAM_KV = auto() + ENGRAM_GATE = auto() DFLASH_ATTN_CONV_BASE = auto() DFLASH_ATTN_CONV_PROJ = auto() DFLASH_FFN_CONV_BASE = auto() @@ -1443,6 +1479,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.DEEPSEEK2OCR: "deepseek2-ocr", MODEL_ARCH.DEEPSEEK32: "deepseek32", MODEL_ARCH.DEEPSEEK4: "deepseek4", + MODEL_ARCH.DEEPSEEK41: "deepseek41", MODEL_ARCH.CHATGLM: "chatglm", MODEL_ARCH.GLM4: "glm4", MODEL_ARCH.GLM4_MOE: "glm4moe", @@ -2143,6 +2180,12 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", + MODEL_TENSOR.DSPARK_MAIN_PROJ: "blk.{bid}.dspark_main_proj", + MODEL_TENSOR.DSPARK_MAIN_NORM: "blk.{bid}.dspark_main_norm", + MODEL_TENSOR.DSPARK_NORM: "blk.{bid}.dspark_norm", + MODEL_TENSOR.ENGRAM_EMBD: "blk.{bid}.engram_embd", + MODEL_TENSOR.ENGRAM_KV: "blk.{bid}.engram_kv", + MODEL_TENSOR.ENGRAM_GATE: "blk.{bid}.engram_gate", MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base", MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj", MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base", @@ -4287,6 +4330,52 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.DEEPSEEK41: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV, + MODEL_TENSOR.ATTN_KV_NORM, + MODEL_TENSOR.ATTN_OUT_A, + MODEL_TENSOR.ATTN_OUT_B, + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + MODEL_TENSOR.ATTN_COMPRESSOR_WKV, + MODEL_TENSOR.ATTN_COMPRESSOR_WGATE, + MODEL_TENSOR.ATTN_COMPRESSOR_NORM, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.ENGRAM_EMBD, + MODEL_TENSOR.ENGRAM_KV, + MODEL_TENSOR.ENGRAM_GATE, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_EXP_PROBS_B_VL, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.DSPARK_MAIN_PROJ, + MODEL_TENSOR.DSPARK_MAIN_NORM, + MODEL_TENSOR.DSPARK_NORM, + MODEL_TENSOR.DSPARK_MARKOV_W1, + MODEL_TENSOR.DSPARK_MARKOV_W2, + MODEL_TENSOR.DSPARK_CONF_PROJ, + ], MODEL_ARCH.ERNIE4_5_MOE: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -6229,6 +6318,7 @@ class GGMLQuantizationType(IntEnum): IQ1_XS = 64 IQ1_XXS = 65 IQ1_XXXS = 66 + MXFP8 = 67 class ExpertGatingFuncType(IntEnum): @@ -6453,6 +6543,7 @@ class VisionProjectorType: GGMLQuantizationType.NVFP4: (64, 4 + 32), GGMLQuantizationType.Q1_0: (128, 2 + 16), GGMLQuantizationType.Q2_0: (64, 2 + 16), + GGMLQuantizationType.MXFP8: (32, 1 + 32), } diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 93ec95eb648e..23ea73aedf00 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1077,6 +1077,75 @@ def add_attention_compress_ratios(self, values: Sequence[int]) -> None: def add_attention_compress_rope_freq_base(self, value: float) -> None: self.add_float32(Keys.Attention.COMPRESS_ROPE_FREQ_BASE.format(arch=self.arch), value) + # deepseek41: cross-layer KV / indexer sharing and two-level candidate selection + def add_attention_kv_source_layers(self, layers: Sequence[int]) -> None: + self.add_key_value(Keys.Attention.KV_SOURCE_LAYERS.format(arch=self.arch), list(layers), GGUFValueType.ARRAY, GGUFValueType.UINT32) + + def add_indexer_source_layers(self, layers: Sequence[int]) -> None: + self.add_key_value(Keys.Attention.Indexer.SOURCE_LAYERS.format(arch=self.arch), list(layers), GGUFValueType.ARRAY, GGUFValueType.UINT32) + + def add_indexer_candidate_source_layer(self, layer: int) -> None: + self.add_int32(Keys.Attention.Indexer.CANDIDATE_SOURCE_LAYER.format(arch=self.arch), layer) + + def add_indexer_candidate_top_k_blocks(self, value: int) -> None: + self.add_uint32(Keys.Attention.Indexer.CANDIDATE_TOP_K_BLOCKS.format(arch=self.arch), value) + + def add_indexer_candidate_block_size(self, value: int) -> None: + self.add_uint32(Keys.Attention.Indexer.CANDIDATE_BLOCK_SIZE.format(arch=self.arch), value) + + # deepseek41 engram + def add_engram_layers(self, layers: Sequence[int]) -> None: + self.add_key_value(Keys.Engram.LAYERS.format(arch=self.arch), list(layers), GGUFValueType.ARRAY, GGUFValueType.UINT32) + + def add_engram_num_embeddings(self, values: Sequence[int]) -> None: + self.add_key_value(Keys.Engram.NUM_EMBEDDINGS.format(arch=self.arch), list(values), GGUFValueType.ARRAY, GGUFValueType.UINT32) + + def add_engram_max_ngram_size(self, value: int) -> None: + self.add_uint32(Keys.Engram.MAX_NGRAM_SIZE.format(arch=self.arch), value) + + def add_engram_vocab_size(self, value: int) -> None: + self.add_uint32(Keys.Engram.VOCAB_SIZE.format(arch=self.arch), value) + + def add_engram_head_count(self, value: int) -> None: + self.add_uint32(Keys.Engram.HEAD_COUNT.format(arch=self.arch), value) + + def add_engram_head_dim(self, value: int) -> None: + self.add_uint32(Keys.Engram.HEAD_DIM.format(arch=self.arch), value) + + def add_engram_pad_token_id(self, value: int) -> None: + self.add_uint32(Keys.Engram.PAD_TOKEN_ID.format(arch=self.arch), value) + + def add_engram_compressed_vocab_size(self, value: int) -> None: + self.add_uint32(Keys.Engram.COMPRESSED_VOCAB_SIZE.format(arch=self.arch), value) + + def add_engram_token_map(self, values: Sequence[int]) -> None: + self.add_key_value(Keys.Engram.TOKEN_MAP.format(arch=self.arch), list(values), GGUFValueType.ARRAY, GGUFValueType.INT32) + + def add_engram_multipliers(self, values: Sequence[int]) -> None: + self.add_key_value(Keys.Engram.MULTIPLIERS.format(arch=self.arch), list(values), GGUFValueType.ARRAY, GGUFValueType.UINT64) + + def add_engram_primes(self, values: Sequence[int]) -> None: + self.add_key_value(Keys.Engram.PRIMES.format(arch=self.arch), list(values), GGUFValueType.ARRAY, GGUFValueType.UINT32) + + # deepseek41 DSpark draft head + def add_dspark_block_size(self, value: int) -> None: + self.add_uint32(Keys.DSpark.BLOCK_SIZE.format(arch=self.arch), value) + + def add_dspark_noise_token_id(self, value: int) -> None: + self.add_uint32(Keys.DSpark.NOISE_TOKEN_ID.format(arch=self.arch), value) + + def add_dspark_target_layers(self, layers: Sequence[int]) -> None: + self.add_key_value(Keys.DSpark.TARGET_LAYERS.format(arch=self.arch), list(layers), GGUFValueType.ARRAY, GGUFValueType.UINT32) + + def add_dspark_markov_rank(self, value: int) -> None: + self.add_uint32(Keys.DSpark.MARKOV_RANK.format(arch=self.arch), value) + + def add_dspark_expert_count(self, value: int) -> None: + self.add_uint32(Keys.DSpark.EXPERT_COUNT.format(arch=self.arch), value) + + def add_dspark_expert_used_count(self, value: int) -> None: + self.add_uint32(Keys.DSpark.EXPERT_USED_COUNT.format(arch=self.arch), value) + def add_hyper_connection_count(self, count: int) -> None: self.add_uint32(Keys.HyperConnection.COUNT.format(arch=self.arch), count) @@ -1092,6 +1161,9 @@ def add_hyper_connection_magnitude(self, value: float) -> None: def add_hyper_connection_low_rank(self, value: int) -> None: self.add_uint32(Keys.HyperConnection.LOW_RANK.format(arch=self.arch), value) + def add_hyper_connection_single_pass(self, value: bool) -> None: + self.add_bool(Keys.HyperConnection.SINGLE_PASS.format(arch=self.arch), value) + def add_ple_layers(self, values: Sequence[int]) -> None: self.add_array(Keys.PerLayerEmbedding.LAYERS.format(arch=self.arch), values) diff --git a/gguf-py/gguf/quants.py b/gguf-py/gguf/quants.py index cc0772913bec..e580846141b3 100644 --- a/gguf-py/gguf/quants.py +++ b/gguf-py/gguf/quants.py @@ -758,6 +758,66 @@ def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: return (d * qs.astype(np.float32)) +class MXFP8(__Quant, qtype=GGMLQuantizationType.MXFP8): + # 32 E4M3 (fn) codes sharing one E8M0 scale: value = e4m3(q) * 2^(e - 127) + # ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf + + @staticmethod + def e4m3_table() -> np.ndarray: + x = np.arange(256, dtype=np.uint32) + exp = ((x >> 3) & 0xF).astype(np.int32) + man = (x & 0x7).astype(np.float64) + v = np.where(exp == 0, man * 2.0**-9, (8.0 + man) * np.exp2(exp - 10.0)) + v = np.where((x & 0x7F) == 0x7F, np.nan, v) + return np.where(x & 0x80, -v, v).astype(np.float32) + + @staticmethod + def e8m0_to_fp32(x: np.ndarray) -> np.ndarray: + # see ggml_e8m0_to_fp32 in ggml-impl.h (0xFF is NaN) + x = x.astype(np.uint32) + bits = np.where(x == 0, np.uint32(0x00400000), x << np.uint32(23)) + return np.where(x == 0xFF, np.float32(np.nan), bits.view(np.float32)) + + @staticmethod + def fp32_to_e4m3(x: np.ndarray) -> np.ndarray: + # see ggml_fp32_to_e4m3 in ggml-impl.h (round to nearest even, saturating) + sign = np.where(np.signbit(x), np.uint8(0x80), np.uint8(0)) + a = np.abs(x).astype(np.float32) + sub = np.rint(a * np.float32(512)).astype(np.int32) + f, e = np.frexp(a) + m = np.rint((f * 2 - 1) * 8).astype(np.int32) + E = e - 1 + 7 + E = np.where(m == 8, E + 1, E) + m = np.where(m == 8, 0, m) + norm = (E << 3) | m + norm = np.where((E > 15) | ((E == 15) & (m == 7)), 0x7E, norm) + code = np.where(a < 2.0**-6, sub, norm) + code = np.where(a >= 448.0, 0x7E, code) + code = np.where(np.isnan(x), 0x7F, code | sign) + return code.astype(np.uint8) + + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + amax = abs(blocks).max(axis=-1, keepdims=True) + + with np.errstate(divide="ignore"): + e = np.where(amax > 0, np.ceil(np.log2(amax / np.float32(448))) + 127, 0) + e = np.clip(e, 0, 254).astype(np.uint8) + + d = cls.e8m0_to_fp32(e) + qs = cls.fp32_to_e4m3(blocks * (np.float32(1) / d)) + + return np.concatenate([e, qs], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + e, qs = np.hsplit(blocks, [1]) + + d = cls.e8m0_to_fp32(e) + + return d * cls.e4m3_table()[qs] + + class NVFP4(__Quant, qtype=GGMLQuantizationType.NVFP4): # E2M1 values doubled (kvalues_mxfp4 convention) kvalues = (0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12) diff --git a/models/templates/deepseek-ai-DeepSeek-V4.1-Flash-1017.jinja b/models/templates/deepseek-ai-DeepSeek-V4.1-Flash-1017.jinja new file mode 100644 index 000000000000..c597891ba11e --- /dev/null +++ b/models/templates/deepseek-ai-DeepSeek-V4.1-Flash-1017.jinja @@ -0,0 +1,275 @@ +{#- DeepSeek-V4.1 chat template (for DeepSeek-V4.1-Flash-1017), ported from encoding/encoding.py in the model repo. + Based on the deepseek-ai/DeepSeek-V4.1-Flash template, with fixes so it matches encoding.py: + tool-result-only user turns do not start a new turn, tool_result content lists keep "[Unsupported x]" blocks, + top-level response_format is attached to the leading system message. + Extra kwargs: thinking_mode ("thinking" | "chat", default "thinking"; enable_thinking / thinking = true/false are + accepted as aliases), drop_thinking (default true), reasoning_effort (int in [1, 100] or "low"/"medium"/"high"/"max", + default "high"). Assistant reasoning is read from reasoning_content, or from thinking as an alias. + Tool call arguments must be objects (string arguments are rendered as a single "arguments" parameter). -#} +{%- set IMG = "<｜deepseek_image｜>" -%} +{%- set DSML = "｜DSML｜" -%} + +{%- if thinking_mode is not defined or thinking_mode is none -%} + {%- if enable_thinking is defined and enable_thinking is not none -%} + {%- set thinking_mode = "thinking" if enable_thinking else "chat" -%} + {%- elif thinking is defined and thinking is not none and thinking is not string -%} + {%- set thinking_mode = "thinking" if thinking else "chat" -%} + {%- else -%} + {%- set thinking_mode = "thinking" -%} + {%- endif -%} +{%- endif -%} +{%- if thinking_mode not in ["chat", "thinking"] -%} + {{- raise_exception("Invalid thinking_mode `" ~ thinking_mode ~ "`, should be \"chat\" or \"thinking\"") -}} +{%- endif -%} +{%- if drop_thinking is not defined or drop_thinking is none -%}{%- set drop_thinking = true -%}{%- endif -%} +{%- if add_generation_prompt is not defined -%}{%- set add_generation_prompt = false -%}{%- endif -%} +{%- if tools is not defined -%}{%- set tools = none -%}{%- endif -%} +{%- if response_format is not defined -%}{%- set response_format = none -%}{%- endif -%} + +{%- if reasoning_effort is not defined or reasoning_effort is none -%}{%- set reasoning_effort = "high" -%}{%- endif -%} +{%- set effort_map = {"low": 50, "medium": 62, "high": 75, "max": 100} -%} +{%- if reasoning_effort is string and reasoning_effort in effort_map -%} + {%- set budget = effort_map[reasoning_effort] -%} +{%- elif reasoning_effort is integer and reasoning_effort >= 1 and reasoning_effort <= 100 -%} + {%- set budget = reasoning_effort -%} +{%- else -%} + {{- raise_exception("Invalid reasoning effort: " ~ reasoning_effort ~ ", should be int within [1,100] or one of \"low\", \"medium\", \"high\", \"max\"") -}} +{%- endif -%} + +{#- Splits/validates "ns::name" against an explicit namespace and returns the qualified name. -#} +{%- macro qualified_name(name, ns) -%} + {%- if ns is mapping -%}{%- set ns = ns.name -%}{%- endif -%} + {%- set parts = name.split("::", 1) -%} + {%- if parts | length == 2 -%} + {%- if ns is not none and ns != parts[0] -%}{{- raise_exception("Conflicting tool namespaces: " ~ ns ~ " != " ~ parts[0]) -}}{%- endif -%} + {%- set ns = parts[0] -%} + {%- set name = parts[1] -%} + {%- endif -%} + {%- if "::" in name -%}{{- raise_exception("Tool name must not contain '::': " ~ name) -}}{%- endif -%} + {%- if ns is not none and "::" in ns -%}{{- raise_exception("Tool namespace must not contain '::': " ~ ns) -}}{%- endif -%} + {{- name if ns is none else ns ~ "::" ~ name -}} +{%- endmacro -%} + +{%- macro render_tools(tool_list) -%} + {%- set ts = namespace(schemas=[]) -%} + {%- for tool in tool_list -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set ns = tool.namespace if tool.namespace is defined and tool.namespace is not none else fn.get("namespace") -%} + {#- JSON assembled by hand (key order kept, name/description replaced in place): llama.cpp has no ** unpacking. -#} + {%- set qname = qualified_name(fn.name, ns) -%} + {%- set sch = namespace(parts=[]) -%} + {%- for k, v in fn.items() if k != "namespace" -%} + {%- if k == "name" -%}{%- set v = qname -%} + {%- elif k == "description" and ns is mapping and ns.description -%}{%- set v = ns.description ~ "\n" ~ (fn.description or "") -%} + {%- endif -%} + {%- set sch.parts = sch.parts + [(k | tojson) ~ ": " ~ (v | tojson)] -%} + {%- endfor -%} + {%- if "description" not in fn and ns is mapping and ns.description -%} + {%- set sch.parts = sch.parts + ['"description": ' ~ ((ns.description ~ "\n") | tojson)] -%} + {%- endif -%} + {%- set ts.schemas = ts.schemas + ["{" ~ sch.parts | join(", ") ~ "}"] -%} + {%- endfor -%} + {{- "## Tools\n\nYou have access to a set of tools to help answer the user's question. You can invoke tools by writing a \"<｜DSML｜ calls>\" block like the following:\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"$TOOL_NAME\">\n<｜DSML｜ parameter name=\"$PARAMETER_NAME\" string=\"true|false\">$PARAMETER_VALUE\n...\n\n<｜DSML｜ invoke name=\"$TOOL_NAME2\">\n...\n\n\n\nString parameters should be specified as is and set `string=\"true\"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string=\"false\"`.\n\nIf thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response.\n\nOtherwise, output directly after with tool calls or final response.\n\n### Available Tool Schemas\n\n" -}} + {{- ts.schemas | join("\n") -}} + {{- "\n\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n" -}} +{%- endmacro -%} + +{#- 1. Merge tool results (and consecutive user turns) into user messages made of content blocks. -#} +{%- set mns = namespace(msgs=[]) -%} +{%- for msg in messages -%} + {%- set content = msg.get("content") -%} + {%- set is_blocks = content is iterable and content is not string and content is not mapping -%} + {%- if content is string and IMG in content -%} + {{- raise_exception("Message content contains image special token " ~ IMG ~ ". Images should be provided as image content blocks.") -}} + {%- endif -%} + {%- set rc = msg.reasoning_content if msg.get("reasoning_content") is not none else msg.get("thinking") -%} + {%- if rc is string and IMG in rc -%} + {{- raise_exception("reasoning_content contains image special token " ~ IMG) -}} + {%- endif -%} + {%- set prev = mns.msgs[-1] if mns.msgs else none -%} + {%- if msg.role == "tool" -%} + {%- set tblock = {"type": "tool_result", "tool_use_id": msg.get("tool_call_id", ""), "content": content if content is not none else ""} -%} + {%- if prev is not none and prev.role == "user" and prev.content_blocks is defined -%} + {%- set mns.msgs = mns.msgs[:-1] + [dict(prev, content_blocks=prev.content_blocks + [tblock])] -%} + {%- else -%} + {%- set mns.msgs = mns.msgs + [{"role": "user", "content_blocks": [tblock]}] -%} + {%- endif -%} + {%- elif msg.role == "user" -%} + {%- set blocks = content if is_blocks else [{"type": "text", "text": content or ""}] -%} + {%- if prev is not none and prev.role == "user" and prev.content_blocks is defined and prev.get("task") is none -%} + {%- set mns.msgs = mns.msgs[:-1] + [dict(prev, content_blocks=prev.content_blocks + blocks)] -%} + {%- else -%} + {%- set mns.msgs = mns.msgs + [dict(msg, content_blocks=blocks)] -%} + {%- endif -%} + {%- else -%} + {%- set mns.msgs = mns.msgs + [msg] -%} + {%- endif -%} +{%- endfor -%} + +{#- Top-level tools / response_format attach to a leading system message, or to a new empty one. -#} +{%- if tools or response_format -%} + {%- if not (mns.msgs and mns.msgs[0].role == "system") -%} + {%- set mns.msgs = [{"role": "system", "content": ""}] + mns.msgs -%} + {%- endif -%} + {%- set first = mns.msgs[0] -%} + {%- if tools and not first.get("tools") -%}{%- set first = dict(first, tools=tools) -%}{%- endif -%} + {%- if response_format and not first.get("response_format") -%}{%- set first = dict(first, response_format=response_format) -%}{%- endif -%} + {%- set mns.msgs = [first] + mns.msgs[1:] -%} +{%- endif -%} + +{#- 2. Reorder tool results to match the tool_calls order of the preceding assistant message. -#} +{%- set sns = namespace(msgs=[], ids=[]) -%} +{%- for m in mns.msgs -%} + {%- if m.role == "assistant" and m.get("tool_calls") -%} + {%- set sns.ids = [] -%} + {%- for tc in m.tool_calls -%} + {%- set sns.ids = sns.ids + [tc.get("id") or (tc.function.get("id", "") if tc.function is defined else "")] -%} + {%- endfor -%} + {%- set sns.msgs = sns.msgs + [m] -%} + {%- elif m.role == "user" and m.content_blocks -%} + {%- set tb = namespace(blocks=[]) -%} + {%- for b in m.content_blocks if b.type == "tool_result" -%} + {%- set order = namespace(idx=0) -%} + {%- for tc_id in sns.ids -%} + {%- if tc_id and tc_id == b.get("tool_use_id", "") -%}{%- set order.idx = loop.index0 -%}{%- endif -%} + {%- endfor -%} + {%- set tb.blocks = tb.blocks + [dict(b, _order=order.idx)] -%} + {%- endfor -%} + {%- if tb.blocks | length > 1 and sns.ids -%} + {%- set sorted_blocks = tb.blocks | sort(attribute="_order") -%} + {%- set rb = namespace(blocks=[], i=0) -%} + {%- for b in m.content_blocks -%} + {%- if b.type == "tool_result" -%} + {%- set rb.blocks = rb.blocks + [sorted_blocks[rb.i]] -%} + {%- set rb.i = rb.i + 1 -%} + {%- else -%} + {%- set rb.blocks = rb.blocks + [b] -%} + {%- endif -%} + {%- endfor -%} + {%- set sns.msgs = sns.msgs + [dict(m, content_blocks=rb.blocks)] -%} + {%- else -%} + {%- set sns.msgs = sns.msgs + [m] -%} + {%- endif -%} + {%- else -%} + {%- set sns.msgs = sns.msgs + [m] -%} + {%- endif -%} +{%- endfor -%} +{%- set msgs = sns.msgs -%} + +{#- Last user turn: users and mid-conversation systems count, except user turns made only of tool results. + Tools anywhere disable drop_thinking. -#} +{%- set state = namespace(last_user=-1, has_tools=false) -%} +{%- for m in msgs -%} + {%- if m.role == "user" or (m.role == "system" and loop.index0 > 0) -%} + {%- set real = namespace(v=true) -%} + {%- if m.get("content_blocks") -%} + {%- set real.v = false -%} + {%- for b in m.content_blocks if b.type in ["text", "image", "image_url"] -%}{%- set real.v = true -%}{%- endfor -%} + {%- endif -%} + {%- if real.v -%}{%- set state.last_user = loop.index0 -%}{%- endif -%} + {%- endif -%} + {%- if m.get("tools") -%}{%- set state.has_tools = true -%}{%- endif -%} +{%- endfor -%} +{%- set effective_drop = drop_thinking and not state.has_tools -%} +{%- set is_thinking = thinking_mode == "thinking" -%} + +{{- "<｜begin▁of▁sentence｜>" -}} +{#- 3. Render each message. -#} +{%- for m in msgs -%} + {%- set idx = loop.index0 -%} + {%- set content = m.get("content") -%} + {%- if content is iterable and content is not string and content is not mapping -%} + {%- set sp = namespace(parts=[]) -%} + {%- for b in content -%} + {%- if b.type == "text" -%}{%- set sp.parts = sp.parts + [b.text or ""] -%} + {%- elif b.type in ["image", "image_url"] -%}{%- set sp.parts = sp.parts + [IMG] -%} + {%- endif -%} + {%- endfor -%} + {%- set content = sp.parts | join("\n\n") -%} + {%- endif -%} + + {%- if idx == 0 -%} + {%- if is_thinking or m.role == "system" -%}{{- "<｜System｜>" -}}{%- endif -%} + {%- if is_thinking -%}{{- "Reasoning Effort: " ~ budget ~ " (range 1-100, the higher the value, the more thorough the reasoning)\n\n" -}}{%- endif -%} + {%- endif -%} + + {%- if m.role == "system" -%} + {%- if idx > 0 -%}{{- "<｜System｜>" -}}{%- endif -%} + {{- content or "" -}} + {%- if m.get("tools") -%}{{- "\n\n" ~ render_tools(m.tools) -}}{%- endif -%} + {%- if m.get("response_format") -%} + {{- "\n\n## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n" ~ (m.response_format | tojson) -}} + {%- endif -%} + + {%- elif m.role == "user" -%} + {{- "<｜User｜>" -}} + {%- set up = namespace(parts=[]) -%} + {%- for b in m.content_blocks -%} + {%- if b.type == "text" -%} + {%- if IMG in (b.text or "") -%}{{- raise_exception("Text block contains image placeholder " ~ IMG ~ ". Images should be separate content blocks.") -}}{%- endif -%} + {%- set up.parts = up.parts + [b.text or ""] -%} + {%- elif b.type in ["image", "image_url"] -%} + {%- set up.parts = up.parts + [IMG] -%} + {%- elif b.type == "tool_result" -%} + {%- set tc = b.get("content", "") -%} + {%- if tc is iterable and tc is not string and tc is not mapping -%} + {%- set sub = namespace(parts=[]) -%} + {%- for s in tc -%} + {%- if s.type == "text" -%} + {%- if IMG in (s.text or "") -%}{{- raise_exception("Text block contains image placeholder " ~ IMG ~ ". Images should be separate content blocks.") -}}{%- endif -%} + {%- set sub.parts = sub.parts + [s.get("text", "")] -%} + {%- elif s.type in ["image", "image_url"] -%}{%- set sub.parts = sub.parts + [IMG] -%} + {%- else -%}{%- set sub.parts = sub.parts + ["[Unsupported " ~ s.type ~ "]"] -%} + {%- endif -%} + {%- endfor -%} + {%- set tc = sub.parts | join("\n\n") -%} + {%- endif -%} + {%- set up.parts = up.parts + ["" ~ tc ~ ""] -%} + {%- else -%} + {%- set up.parts = up.parts + ["[Unsupported " ~ b.type ~ "]"] -%} + {%- endif -%} + {%- endfor -%} + {{- up.parts | join("\n\n") -}} + + {%- elif m.role == "latest_reminder" -%} + {{- "<｜latest_reminder｜>" ~ (content or "") -}} + + {%- elif m.role == "assistant" -%} + {%- set prev_has_task = idx > 0 and msgs[idx - 1].get("task") is not none -%} + {%- if is_thinking and not prev_has_task and (not effective_drop or idx > state.last_user) -%} + {{- ((m.reasoning_content if m.get("reasoning_content") is not none else m.get("thinking")) or "") ~ "" -}} + {%- endif -%} + {{- content or "" -}} + {%- if m.get("tool_calls") -%} + {%- set tcs = namespace(lines=[]) -%} + {%- for tc in m.tool_calls -%} + {%- set fn = tc.function if tc.function is defined else tc -%} + {%- set params = namespace(lines=[]) -%} + {%- set args = fn.arguments if fn.arguments is mapping else {"arguments": fn.arguments} -%} + {%- for k, v in args.items() -%} + {%- set params.lines = params.lines + ["<" ~ DSML ~ " parameter name=\"" ~ k ~ "\" string=\"" ~ ("true" if v is string else "false") ~ "\">" ~ (v if v is string else v | tojson) ~ ""] -%} + {%- endfor -%} + {%- set name = qualified_name(fn.name, tc.get("namespace") or fn.get("namespace")) -%} + {%- set tcs.lines = tcs.lines + ["<" ~ DSML ~ " invoke name=\"" ~ name ~ "\">\n" ~ params.lines | join("\n") ~ "\n"] -%} + {%- endfor -%} + {{- "\n\n<" ~ DSML ~ " calls>\n" ~ tcs.lines | join("\n") ~ "\n" -}} + {%- endif -%} + {%- if not m.get("wo_eos", false) -%}{{- "<｜end▁of▁sentence｜>" -}}{%- endif -%} + + {%- else -%} + {{- raise_exception("Unknown role: " ~ m.role) -}} + {%- endif -%} + + {#- Transition tokens, only when the next message is an assistant/latest_reminder turn or this is the last one. -#} + {%- if loop.last or msgs[idx + 1].role in ["assistant", "latest_reminder"] -%} + {%- set task = m.get("task") -%} + {%- if task is not none -%} + {%- set task_tokens = {"action": "<｜action｜>", "query": "<｜query｜>", "authority": "<｜authority｜>", "domain": "<｜domain｜>", "title": "<｜title｜>", "read_url": "<｜read_url｜>"} -%} + {%- if task not in task_tokens -%}{{- raise_exception("Invalid task: '" ~ task ~ "'") -}}{%- endif -%} + {%- if task == "action" -%}{{- "<｜Assistant｜>" ~ ("" if is_thinking else "") -}}{%- endif -%} + {{- task_tokens[task] -}} + {%- elif (m.role == "user" or (m.role == "system" and idx > 0)) and (add_generation_prompt or not loop.last) -%} + {{- "<｜Assistant｜>" ~ ("" if is_thinking and (not effective_drop or idx >= state.last_user) else "") -}} + {%- endif -%} + {%- endif -%} +{%- endfor -%} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index ffa35bdea7d4..d17a8a6d2a47 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -85,6 +85,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" }, { LLM_ARCH_DEEPSEEK32, "deepseek32" }, { LLM_ARCH_DEEPSEEK4, "deepseek4" }, + { LLM_ARCH_DEEPSEEK41, "deepseek41" }, { LLM_ARCH_CHATGLM, "chatglm" }, { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, @@ -309,6 +310,11 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" }, { LLM_KV_ATTENTION_INDEXER_KPOOL, "%s.attention.indexer.kpool" }, { LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, "%s.attention.indexer.kpool_select_tail" }, + { LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, "%s.attention.indexer.source_layers" }, + { LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, "%s.attention.indexer.candidate_source_layer" }, + { LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, "%s.attention.indexer.candidate_top_k_blocks" }, + { LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, "%s.attention.indexer.candidate_block_size" }, + { LLM_KV_ATTENTION_KV_SOURCE_LAYERS, "%s.attention.kv_source_layers" }, { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" }, { LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" }, { LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" }, @@ -321,6 +327,26 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_HYPER_CONNECTION_EPSILON, "%s.hyper_connection.epsilon" }, { LLM_KV_HYPER_CONNECTION_MAGNITUDE, "%s.hyper_connection.magnitude" }, { LLM_KV_HYPER_CONNECTION_LOW_RANK, "%s.hyper_connection.low_rank" }, + { LLM_KV_HYPER_CONNECTION_SINGLE_PASS, "%s.hyper_connection.single_pass" }, + + { LLM_KV_ENGRAM_LAYERS, "%s.engram.layers" }, + { LLM_KV_ENGRAM_NUM_EMBEDDINGS, "%s.engram.num_embeddings" }, + { LLM_KV_ENGRAM_MAX_NGRAM_SIZE, "%s.engram.max_ngram_size" }, + { LLM_KV_ENGRAM_VOCAB_SIZE, "%s.engram.vocab_size" }, + { LLM_KV_ENGRAM_HEAD_COUNT, "%s.engram.head_count" }, + { LLM_KV_ENGRAM_HEAD_DIM, "%s.engram.head_dim" }, + { LLM_KV_ENGRAM_PAD_TOKEN_ID, "%s.engram.pad_token_id" }, + { LLM_KV_ENGRAM_COMPRESSED_VOCAB_SIZE, "%s.engram.compressed_vocab_size" }, + { LLM_KV_ENGRAM_TOKEN_MAP, "%s.engram.token_map" }, + { LLM_KV_ENGRAM_MULTIPLIERS, "%s.engram.multipliers" }, + { LLM_KV_ENGRAM_PRIMES, "%s.engram.primes" }, + + { LLM_KV_DSPARK_BLOCK_SIZE, "%s.dspark.block_size" }, + { LLM_KV_DSPARK_NOISE_TOKEN_ID, "%s.dspark.noise_token_id" }, + { LLM_KV_DSPARK_TARGET_LAYERS, "%s.dspark.target_layers" }, + { LLM_KV_DSPARK_MARKOV_RANK, "%s.dspark.markov_rank" }, + { LLM_KV_DSPARK_EXPERT_COUNT, "%s.dspark.expert_count" }, + { LLM_KV_DSPARK_EXPERT_USED_COUNT, "%s.dspark.expert_used_count" }, { LLM_KV_PLE_LAYERS, "%s.ple.layers" }, { LLM_KV_PLE_NGRAM_SIZE, "%s.ple.ngram_size" }, @@ -775,6 +801,12 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, + { LLM_TENSOR_DSPARK_MAIN_PROJ, "blk.%d.dspark_main_proj" }, + { LLM_TENSOR_DSPARK_MAIN_NORM, "blk.%d.dspark_main_norm" }, + { LLM_TENSOR_DSPARK_NORM, "blk.%d.dspark_norm" }, + { LLM_TENSOR_ENGRAM_EMBD, "blk.%d.engram_embd" }, + { LLM_TENSOR_ENGRAM_KV, "blk.%d.engram_kv" }, + { LLM_TENSOR_ENGRAM_GATE, "blk.%d.engram_gate" }, { LLM_TENSOR_DFLASH_ATTN_CONV_BASE, "blk.%d.attn_conv_base" }, { LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "blk.%d.attn_conv_proj" }, { LLM_TENSOR_DFLASH_FFN_CONV_BASE, "blk.%d.ffn_conv_base" }, @@ -1113,6 +1145,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, {LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DSPARK_MAIN_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DSPARK_MAIN_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_DSPARK_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + // engram tables are gathered rows, kept host-side (never repacked / offloaded) + {LLM_TENSOR_ENGRAM_EMBD, {LLM_TENSOR_LAYER_INPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_ENGRAM_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ENGRAM_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_DFLASH_ATTN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_DFLASH_FFN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, @@ -1232,6 +1271,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_MINIMAX_01: return true; default: @@ -1258,6 +1298,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: @@ -1275,6 +1316,7 @@ bool llm_arch_supports_mixed_batch(const llm_arch & arch) { switch (arch) { case LLM_ARCH_COGVLM: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_GRANITE_SWITCH: case LLM_ARCH_EAGLE3: case LLM_ARCH_DFLASH: diff --git a/src/llama-arch.h b/src/llama-arch.h index 14ce29ab3238..94b28b3429f1 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -90,6 +90,7 @@ enum llm_arch { LLM_ARCH_DEEPSEEK2OCR, LLM_ARCH_DEEPSEEK32, LLM_ARCH_DEEPSEEK4, + LLM_ARCH_DEEPSEEK41, LLM_ARCH_CHATGLM, LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, @@ -314,6 +315,11 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_TYPES, LLM_KV_ATTENTION_INDEXER_KPOOL, LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, + LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, + LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, + LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, + LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, + LLM_KV_ATTENTION_KV_SOURCE_LAYERS, LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, LLM_KV_ATTENTION_OUTPUT_LORA_RANK, LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, @@ -326,6 +332,26 @@ enum llm_kv { LLM_KV_HYPER_CONNECTION_EPSILON, LLM_KV_HYPER_CONNECTION_MAGNITUDE, LLM_KV_HYPER_CONNECTION_LOW_RANK, + LLM_KV_HYPER_CONNECTION_SINGLE_PASS, + + LLM_KV_ENGRAM_LAYERS, + LLM_KV_ENGRAM_NUM_EMBEDDINGS, + LLM_KV_ENGRAM_MAX_NGRAM_SIZE, + LLM_KV_ENGRAM_VOCAB_SIZE, + LLM_KV_ENGRAM_HEAD_COUNT, + LLM_KV_ENGRAM_HEAD_DIM, + LLM_KV_ENGRAM_PAD_TOKEN_ID, + LLM_KV_ENGRAM_COMPRESSED_VOCAB_SIZE, + LLM_KV_ENGRAM_TOKEN_MAP, + LLM_KV_ENGRAM_MULTIPLIERS, + LLM_KV_ENGRAM_PRIMES, + + LLM_KV_DSPARK_BLOCK_SIZE, + LLM_KV_DSPARK_NOISE_TOKEN_ID, + LLM_KV_DSPARK_TARGET_LAYERS, + LLM_KV_DSPARK_MARKOV_RANK, + LLM_KV_DSPARK_EXPERT_COUNT, + LLM_KV_DSPARK_EXPERT_USED_COUNT, LLM_KV_PLE_LAYERS, LLM_KV_PLE_NGRAM_SIZE, @@ -781,6 +807,12 @@ enum llm_tensor { LLM_TENSOR_DSPARK_MARKOV_W1, LLM_TENSOR_DSPARK_MARKOV_W2, LLM_TENSOR_DSPARK_CONF_PROJ, + LLM_TENSOR_DSPARK_MAIN_PROJ, + LLM_TENSOR_DSPARK_MAIN_NORM, + LLM_TENSOR_DSPARK_NORM, + LLM_TENSOR_ENGRAM_EMBD, + LLM_TENSOR_ENGRAM_KV, + LLM_TENSOR_ENGRAM_GATE, LLM_TENSOR_DFLASH_ATTN_CONV_BASE, LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, LLM_TENSOR_DFLASH_FFN_CONV_BASE, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9ee64b566779..442b1df8b64c 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1323,8 +1323,11 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void void llama_context::set_embeddings(bool value) { LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); - // a pooled context computes the logits only without embeddings, so the toggle changes the graph shape - if (cparams.embeddings != value && cparams.pooling_type != LLAMA_POOLING_TYPE_NONE) { + // a pooled context computes the logits only without embeddings, so the toggle changes the graph shape. + // with pooling NONE the reserved graph (all tokens output) can also differ from the generation graph: + // llama-server --embeddings --pooling none segfaulted in ggml_gallocr on the first /completion (any arch), + // so re-reserve on every toggle + if (cparams.embeddings != value) { sched_need_reserve = true; } @@ -2520,6 +2523,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_CLEF || model.arch == LLM_ARCH_QWEN4EXP || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_DEEPSEEK41 || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || @@ -4055,7 +4059,7 @@ llama_context * llama_init_from_model( } } - if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) { + if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) && params.type_k != params.type_v) { LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v)); return nullptr; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 2bc3a291a456..9106f011f9ce 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1968,7 +1968,7 @@ ggml_tensor * llm_graph_context::build_ffn( const float limit = hparams.swiglu_clamp_shexp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit); } else { tmp = ggml_clamp(ctx0, tmp, -limit, limit); @@ -2187,6 +2187,12 @@ ggml_tensor * llm_graph_context::build_moe_ffn( if (gating_op == LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { ggml_prec_set_acc(logits, GGML_PREC_F32); } + if (arch == LLM_ARCH_DEEPSEEK41 && !ggml_is_quantized(gate_inp->type)) { + // F32 router: keep the logits off TF32 on CUDA (top-6 of 384 is near-tie sensitive; with TF32 the routing + // depended on the ubatch size, 4..16-token batches being the worst) + // skipped for quantized routers (third-party GGUFs): a quantized src0 with PREC_F32 asserts in MMQ + ggml_prec_set_src(logits, GGML_PREC_F32, 1); + } cb(logits, "ffn_moe_logits", il); } else { logits = probs_in; @@ -2298,8 +2304,13 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * weights_sum = ggml_sum_rows(ctx0, weights); // [1, n_tokens] cb(weights_sum, "ffn_moe_weights_sum", il); - // Avoid division by zero, clamp to smallest number representable by F16 - weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5, INFINITY); + if (arch == LLM_ARCH_DEEPSEEK41) { + // reference: weights / (sum + 1e-20), not the F16-floor clamp (model.py Gate.forward) + weights_sum = ggml_scale_bias(ctx0, weights_sum, 1.0f, 1e-20f); + } else { + // Avoid division by zero, clamp to smallest number representable by F16 + weights_sum = ggml_clamp(ctx0, weights_sum, 6.103515625e-5, INFINITY); + } cb(weights_sum, "ffn_moe_weights_sum_clamped", il); weights = ggml_div(ctx0, weights, weights_sum); // [n_expert_used, n_tokens] @@ -2389,7 +2400,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( const float limit = hparams.swiglu_clamp_exp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) { + if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) { cur = ggml_swiglu_clamp(ctx0, cur, up, limit); } else { up = ggml_clamp(ctx0, up, -limit, limit); @@ -2971,6 +2982,10 @@ ggml_tensor * llm_graph_context::build_attn_mha( // note: this op tends to require high floating point range // while for some models F16 is enough, for others it is not, so we default to F32 here ggml_prec_set_acc(kq, GGML_PREC_F32); + if (arch == LLM_ARCH_DEEPSEEK41) { + // with an F32 KV cache: no TF32 rounding of K / Q on CUDA (no-op for F16 / quantized caches) + ggml_prec_set_src(kq, GGML_PREC_F32, 1); + } if (arch == LLM_ARCH_GROK) { // need to do the following: @@ -3010,6 +3025,9 @@ ggml_tensor * llm_graph_context::build_attn_mha( } ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq); + if (arch == LLM_ARCH_DEEPSEEK41) { + ggml_prec_set_src(kqv, GGML_PREC_F32, 1); + } cb(kqv, "kqv", il); // for MLA with the absorption optimization, we need to "decompress" from MQA back to MHA diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 719d7081dd21..875bcb797571 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -326,12 +326,22 @@ struct llama_hparams { uint32_t dsv4_o_group_count = 0; uint32_t dsv4_o_lora_rank = 0; uint32_t dsv4_hc_mult = 0; + bool dsv4_hc_single_pass = false; // DeepSeek-V4.1 DSpark draft: lagged pre-mix, no hc_head uint32_t dsv4_hc_sinkhorn_iters = 0; uint32_t dsv4_hash_layer_count = 0; float dsv4_compress_rope_base = 0.0f; float dsv4_hc_eps = 0.0f; std::array dsv4_compress_ratios; + // deepseek41: cross-layer sharing of the compressed KV / index K and of the indexer top-k + std::array dsv41_kv_src; // layer whose compressed KV (and index K) this layer reads, -1 = none + std::array dsv41_topk_src; // index-source layer whose top-k this layer attends with, -1 = none + int32_t dsv41_cand_src = -1; // candidate_source_layer, -1 = two-level selection off + bool dsv41_tok_hist = false; // Engram: KV cells keep (and serialize) token ids for the n-gram history + uint32_t dsv41_hist_len = 0; // Engram: number of predecessor tokens an n-gram hash reads (max_ngram - 1) + uint32_t dsv41_cand_topk_blocks = 0; + uint32_t dsv41_cand_block_size = 0; + // 0 = full rank (DeepSeek-V4) uint32_t hc_low_rank = 0; diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index f8caae9d6be5..c04c9bd58073 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -432,6 +432,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( const llama_ubatch & ubatch, uint32_t ratio, bool overlap, + bool pad_blocks, uint32_t state_size, uint32_t kv_size, uint32_t n_stream, @@ -548,7 +549,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( } } - if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) { + if (pad_blocks && !plan.state_pos.empty()) { assert(kv_size > 0); // Pad each stream to the reserve plan's block count. @@ -603,7 +604,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( } } - if (ratio == DSV4_HCA_RATIO && !plan.state_pos.empty() && plan.state_write_idxs.empty()) { + if (!pad_blocks && !plan.state_pos.empty() && plan.state_write_idxs.empty()) { assert(kv_size > 0); // the last slot must not be live, or the dummy write would corrupt it; // a full stream implies a completed block, which implies real writes @@ -744,6 +745,7 @@ static std::vector dsv4_build_comp_plans const std::vector & ubatches, uint32_t ratio, bool overlap, + bool pad_blocks, uint32_t state_size, uint32_t kv_size, uint32_t n_stream, @@ -755,7 +757,7 @@ static std::vector dsv4_build_comp_plans // the first ubatch touching a seq consumes its rollback restore std::vector rs(rs_idx); for (const llama_ubatch & ubatch : ubatches) { - plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs)); + plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, pad_blocks, state_size, kv_size, n_stream, n_rs_seq, rs)); for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { const llama_seq_id seq_id = ubatch.seq_id_unq[s]; @@ -1056,6 +1058,10 @@ uint32_t llama_dsv4_comp_state::get_ratio() const { return ratio; } +bool llama_dsv4_comp_state::is_overlap() const { + return state_size == 2*ratio; +} + uint32_t llama_dsv4_comp_state::get_state_size() const { return state_size; } @@ -1261,6 +1267,25 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( v_trans, offload, swa_full, unified_raw, kv_size, n_seq_max, n_ubatch, n_pad, nullptr, filter_raw, reuse, nullptr); + // V4: CSA = overlapped ratio-4 pooling (+ its own indexer compressor, "lid"), HCA = ratio-128 pooling. + // V4.1: CSA slot = ratio-2 pooling, HCA slot = ratio-1 projection, no overlap, no "lid": the index K is derived from + // the compressed latent and stored next to it in the same row ([kv | index K]), and only the KV source layers own + // rows (consumers read the source layer's cache). + const bool is_v41 = model.arch == LLM_ARCH_DEEPSEEK41; + const uint32_t ratio_csa = is_v41 ? 2u : DSV4_CSA_RATIO; + const uint32_t ratio_hca = is_v41 ? 1u : DSV4_HCA_RATIO; + + if (is_v41) { + const uint32_t n_embd_comp = model.hparams.n_embd_head_k() + model.hparams.indexer_head_size; + for (llama_hparams * hp : { &hparams_csa, &hparams_hca }) { + std::fill(hp->n_head_kv_arr.begin(), hp->n_head_kv_arr.end(), 1); + hp->n_embd_head_k_full = n_embd_comp; + hp->n_embd_head_v_full = n_embd_comp; + hp->n_embd_head_k_swa = n_embd_comp; + hp->n_embd_head_v_swa = n_embd_comp; + } + } + dsv4_make_k_only(hparams_csa); dsv4_make_k_only(hparams_hca); @@ -1272,12 +1297,19 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( hparams_lid.rope_type = LLAMA_ROPE_TYPE_NEOX; dsv4_make_k_only(hparams_lid); + const auto is_comp_owner = [&](int32_t il, uint32_t ratio) { + if (model.hparams.dsv4_compress_ratios[il] != ratio) { + return false; + } + return !is_v41 || model.hparams.dsv41_kv_src[il] == il; + }; + const layer_filter_cb filter_csa = [&](int32_t il) { if (filter && !filter(il)) { return false; } - return model.hparams.dsv4_compress_ratios[il] == DSV4_CSA_RATIO; + return is_comp_owner(il, ratio_csa); }; const layer_filter_cb filter_hca = [&](int32_t il) { @@ -1285,52 +1317,60 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( return false; } - return model.hparams.dsv4_compress_ratios[il] == DSV4_HCA_RATIO; + return is_comp_owner(il, ratio_hca); + }; + + const layer_filter_cb filter_lid = [&](int32_t il) { + return !is_v41 && filter_csa(il); }; const bool unified_compressed = false; - LLAMA_LOG_INFO("%s: creating DSV4 CSA compressed KV cache, size = %u cells\n", - __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO)); + LLAMA_LOG_INFO("%s: creating DSV4 CSA compressed KV cache, ratio = %u, size = %u cells\n", + __func__, ratio_csa, dsv4_comp_size(kv_size, ratio_csa)); kv_csa = std::make_unique( model, hparams_csa, type_k, type_v, - v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad, + v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr); - LLAMA_LOG_INFO("%s: creating DSV4 HCA compressed KV cache, size = %u cells\n", - __func__, dsv4_comp_size(kv_size, DSV4_HCA_RATIO)); + LLAMA_LOG_INFO("%s: creating DSV4 HCA compressed KV cache, ratio = %u, size = %u cells\n", + __func__, ratio_hca, dsv4_comp_size(kv_size, ratio_hca)); kv_hca = std::make_unique( model, hparams_hca, type_k, type_v, - v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_HCA_RATIO), 256u), n_seq_max, n_pad, + v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_hca), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_hca, nullptr, nullptr); LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer KV cache, size = %u cells\n", - __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO)); + __func__, dsv4_comp_size(kv_size, ratio_csa)); kv_lid = std::make_unique( model, hparams_lid, type_k, type_v, - v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad, - 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr); + v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, ratio_csa), 256u), n_seq_max, n_pad, + 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_lid, nullptr, nullptr); LLAMA_LOG_INFO("%s: creating DSV4 CSA compressor state\n", __func__); + // overlapped pooling keeps 2*ratio rows of 2x width; plain pooling keeps ratio rows of head width + const uint32_t csa_state_size = is_v41 ? ratio_csa : 2*ratio_csa; + const uint32_t csa_state_embd = is_v41 ? model.hparams.n_embd_head_k() : 2*model.hparams.n_embd_head_k(); + csa_state = std::make_unique( - model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, - 2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa); + model, offload, unified_compressed, n_seq_max, ratio_csa, csa_state_size, + csa_state_embd, n_rs_seq, "csa", filter_csa); LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__); hca_state = std::make_unique( - model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO, + model, offload, unified_compressed, n_seq_max, ratio_hca, ratio_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( model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, - 2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa); + 2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_lid); // DSV4 attention reads compressed-K / compressor-state rows that the current // graph does not necessarily overwrite; uninitialized buffer contents would @@ -1339,12 +1379,78 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( clear_compressed(-1, true); } +// DeepSeek-V4.1 Engram hashes the n-gram ending at each token, so a token that belongs to several sequences +// (llama-batched shares the prompt that way) has one well-defined input only if all of its sequences agree on +// its predecessors. Predecessors come from the same batch first, then from the cached cells. Returns false (and +// logs) for the first ambiguous token. +static bool dsv41_shared_history_ok(const llama_batch & batch, const llama_kv_cache * kv, uint32_t n_hist) { + if (n_hist == 0 || !batch.token || !batch.n_seq_id) { + return true; + } + + bool any_shared = false; + for (int32_t i = 0; i < batch.n_tokens && !any_shared; ++i) { + any_shared = batch.n_seq_id[i] > 1; + } + if (!any_shared) { + return true; + } + + std::map, llama_token> in_batch; + for (int32_t i = 0; i < batch.n_tokens; ++i) { + for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { + in_batch[{ batch.seq_id[i][s], batch.pos[i] }] = batch.token[i]; + } + } + + const auto tok_at = [&](llama_seq_id seq_id, llama_pos p) -> llama_token { + // largest position <= p of this sequence, batch first (the batch continues the cache) + auto it = in_batch.upper_bound({ seq_id, p }); + if (it != in_batch.begin()) { + --it; + if (it->first.first == seq_id) { + return it->second; + } + } + return kv ? kv->seq_pos_tok_le(seq_id, p) : LLAMA_TOKEN_NULL; + }; + + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] <= 1) { + continue; + } + for (uint32_t d = 1; d <= n_hist; ++d) { + const llama_pos p = batch.pos[i] - (llama_pos) d; + if (p < 0) { + break; + } + const llama_token t0 = tok_at(batch.seq_id[i][0], p); + for (int32_t s = 1; s < batch.n_seq_id[i]; ++s) { + const llama_token ts = tok_at(batch.seq_id[i][s], p); + if (ts != t0) { + LLAMA_LOG_ERROR("%s: token %d at pos %d is shared by seq %d and seq %d, whose Engram n-gram histories " + "differ at pos %d (%d vs %d); decode such tokens per sequence\n", __func__, + i, batch.pos[i], batch.seq_id[i][0], batch.seq_id[i][s], p, t0, ts); + return false; + } + } + } + } + + return true; +} + llama_memory_context_ptr llama_kv_cache_dsv4::init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { GGML_UNUSED(embd_all); + if (hparams_raw.dsv41_tok_hist && + !dsv41_shared_history_ok(balloc.get_batch(), kv_raw->get_swa(), hparams_raw.dsv41_hist_len)) { + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + const bool raw_per_seq = kv_raw->get_base()->get_n_stream() != 1; const bool comp_per_seq = csa_state->get_n_stream() > 1; const bool has_coupled = dsv4_batch_has_coupled(balloc.get_batch()); @@ -1475,9 +1581,9 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 bool res = true; res = res & kv_raw->seq_rm(seq_id, p0, -1); - res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1); - res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1); - res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1); + res = res & kv_csa->seq_rm(seq_id, p0/csa_state->get_ratio(), -1); + res = res & kv_hca->seq_rm(seq_id, p0/hca_state->get_ratio(), -1); + res = res & kv_lid->seq_rm(seq_id, p0/lid_state->get_ratio(), -1); return res; } @@ -1613,11 +1719,11 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id //FIXME : note that we conflate token positions with rows, which is not true for multi-modal case. const uint32_t n_rows_csa = seq_id >= 0 ? - dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_csa->get_size()) : kv_csa->get_size(); + dsv4_state_n_used_k_rows(pos_max, csa_state->get_ratio(), kv_csa->get_size()) : kv_csa->get_size(); const uint32_t n_rows_hca = seq_id >= 0 ? - dsv4_state_n_used_k_rows(pos_max, DSV4_HCA_RATIO, kv_hca->get_size()) : kv_hca->get_size(); + dsv4_state_n_used_k_rows(pos_max, hca_state->get_ratio(), kv_hca->get_size()) : kv_hca->get_size(); const uint32_t n_rows_lid = seq_id >= 0 ? - dsv4_state_n_used_k_rows(pos_max, DSV4_CSA_RATIO, kv_lid->get_size()) : kv_lid->get_size(); + dsv4_state_n_used_k_rows(pos_max, lid_state->get_ratio(), kv_lid->get_size()) : kv_lid->get_size(); dsv4_state_write_k_cache(io, kv_csa.get(), seq_id, flags, n_rows_csa); dsv4_state_write_k_cache(io, kv_hca.get(), seq_id, flags, n_rows_hca); @@ -1932,6 +2038,11 @@ void llama_kv_cache_dsv4_raw_context::set_input_kq_mask(ggml_tensor * dst, const kv_swa->set_input_kq_mask(dst, ubatch, causal_attn); } +void llama_kv_cache_dsv4_raw_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const { + GGML_ASSERT(kv_swa != nullptr); + kv_swa->get_prev_tokens(ubatch, n, res); +} + void llama_kv_cache_dsv4_raw_context::set_input_k_rot(ggml_tensor * dst) const { kv_swa->set_input_k_rot(dst); } @@ -2052,13 +2163,13 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( std::vector ubatches, std::vector ubatches_raw) : ubatches(std::move(ubatches)), - plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true, + plans_csa(dsv4_build_comp_plans(this->ubatches, kv->get_csa_state()->get_ratio(), kv->get_csa_state()->is_overlap(), true, 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, + plans_hca(dsv4_build_comp_plans(this->ubatches, kv->get_hca_state()->get_ratio(), kv->get_hca_state()->is_overlap(), false, 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, + plans_lid(dsv4_build_comp_plans(this->ubatches, kv->get_lid_state()->get_ratio(), kv->get_lid_state()->is_overlap(), 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( @@ -2222,7 +2333,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ } reserve_plan_csa = dsv4_build_reserve_comp_plan( - ubatch, DSV4_CSA_RATIO, true, + ubatch, csa_state->get_ratio(), csa_state->is_overlap(), 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; @@ -2236,7 +2347,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ } reserve_plan_hca = dsv4_build_reserve_comp_plan( - ubatch, DSV4_HCA_RATIO, false, + ubatch, hca_state->get_ratio(), hca_state->is_overlap(), 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; @@ -2250,7 +2361,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ } reserve_plan_lid = dsv4_build_reserve_comp_plan( - ubatch, DSV4_CSA_RATIO, true, + ubatch, lid_state->get_ratio(), lid_state->is_overlap(), 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; diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index ce39867c0342..2a6e14a0da8e 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -31,6 +31,7 @@ class llama_dsv4_comp_state { void apply_copies(const stream_copy_info & sc_info) const; uint32_t get_ratio() const; + bool is_overlap() const; // overlapped (V4 CSA) pooling: state holds 2*ratio rows of 2x width uint32_t get_state_size() const; uint32_t get_n_stream() const; uint32_t get_n_rs_seq() const; @@ -214,6 +215,9 @@ class llama_kv_cache_dsv4_raw_context : public llama_memory_context_i { void set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_k_rot(ggml_tensor * dst) const; + // token history of the raw (window) cells, for host-side n-gram hashing (deepseek41 Engram) + void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const; + private: size_t i_next = 0; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index b7dfadef4638..77651694677e 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1150,7 +1150,8 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & cells.pos_set(idx, ubatch.pos[i]); - if (ubatch.is_pos_2d() || ubatch.token || hparams.ple_n_heads > 0) { + // dsv41_tok_hist: embd cells must record LLAMA_TOKEN_NULL (Engram treats it as DEAD), never a stale id + if (ubatch.is_pos_2d() || ubatch.token || hparams.ple_n_heads > 0 || hparams.dsv41_tok_hist) { llama_kv_cell_ext ext; if (ubatch.is_pos_2d()) { @@ -1846,8 +1847,8 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const { } bool llama_kv_cache::has_cell_ext() const { - // M-RoPE needs the 2D position, the PLE n-gram hash needs the token id - return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0; + // M-RoPE needs the 2D position, the PLE / Engram n-gram hashes need the token id + return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0 || hparams.dsv41_tok_hist; } void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const { @@ -1909,6 +1910,14 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st } } +llama_token llama_kv_cache::seq_pos_tok_le(llama_seq_id seq_id, llama_pos p) const { + if (!has_cell_ext() || seq_id < 0 || (size_t) seq_id >= seq_to_stream.size() || p < 0) { + return LLAMA_TOKEN_NULL; + } + + return v_cells[seq_to_stream[seq_id]].seq_pos_tok_le(seq_id, p); +} + size_t llama_kv_cache::total_size() const { size_t size = 0; diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 9bfb8f231fe6..9d19f772ee65 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -250,6 +250,9 @@ class llama_kv_cache : public llama_memory_i { // note: used by n-gram input embeddings void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const; + // the token of sequence seq_id at the largest cached position <= p (LLAMA_TOKEN_NULL if none or no cell ext) + llama_token seq_pos_tok_le(llama_seq_id seq_id, llama_pos p) const; + private: const llama_model & model; const llama_hparams & hparams; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index bc8faa4e0e96..6c14e9dbe5f4 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -417,6 +417,7 @@ namespace GGUFMeta { template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); @@ -444,6 +445,7 @@ namespace GGUFMeta { template bool llama_model_loader::get_key (enum llm_kv kid, bool & result, bool required); template bool llama_model_loader::get_key (enum llm_kv kid, float & result, bool required); template bool llama_model_loader::get_key (enum llm_kv kid, uint32_t & result, bool required); + template bool llama_model_loader::get_key (enum llm_kv kid, int32_t & result, bool required); template bool llama_model_loader::get_key(enum llm_kv kid, std::string & result, bool required); template<> @@ -1206,7 +1208,8 @@ struct ggml_tensor * llama_model_loader::create_tensor( // sanity checks if (info.layer == LLM_TENSOR_LAYER_INPUT || info.layer == LLM_TENSOR_LAYER_OUTPUT) { - if (tn.bid != -1) { + // deepseek41 Engram tables are per-layer but host-resident gather tables, placed like token_embd + if (tn.bid != -1 && tn.tensor != LLM_TENSOR_ENGRAM_EMBD) { GGML_ABORT("input/output layer tensor %s used with a layer number", tn.str().c_str()); } } else { @@ -1273,6 +1276,15 @@ struct ggml_tensor * llama_model_loader::create_tensor( } } + // deepseek41 Engram tables (~100 GB each on the full model) are host-resident gather tables: the graph only + // get_rows a few rows per token on the CPU. Never let an override or a split mode put them in device memory. + if (tn.tensor == LLM_TENSOR_ENGRAM_EMBD && !ggml_backend_buft_is_host(buft)) { + LLAMA_LOG_WARN("llama_model_loader: %s (%zu MiB %s) must stay in host memory; ignoring placement on %s, using %s\n", + tn.str().c_str(), ggml_nbytes(t_meta) / 1024 / 1024, ggml_type_name(t_meta->type), + ggml_backend_buft_name(buft), ggml_backend_buft_name(ggml_backend_cpu_buffer_type())); + buft = ggml_backend_cpu_buffer_type(); + } + // avoid using a host buffer when using mmap auto * buft_dev = ggml_backend_buft_get_device(buft); if (use_mmap && buft_dev && buft == ggml_backend_dev_host_buffer_type(buft_dev)) { diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 46b1d2ae2744..b69832cf0979 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -8,6 +8,7 @@ #include "llama-hparams.h" #include "llama-model.h" #include "llama-vocab.h" +#include "models/models.h" #include #include @@ -338,6 +339,49 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count); add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); + if (model->arch == LLM_ARCH_DEEPSEEK41) { + const auto * m41 = dynamic_cast(model); + GGML_ASSERT(m41 != nullptr); + std::vector kv_src; + std::vector idx_src; + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + if (hparams.dsv41_kv_src[il] == (int32_t) il) { + kv_src.push_back(il); + } + if (hparams.dsv41_topk_src[il] == (int32_t) il) { + idx_src.push_back(il); + } + } + add_kv(LLM_KV_ATTENTION_KV_SOURCE_LAYERS, kv_src); + add_kv(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, idx_src); + if (hparams.dsv41_cand_src >= 0) { + add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, hparams.dsv41_cand_src); + add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, hparams.dsv41_cand_topk_blocks); + add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, hparams.dsv41_cand_block_size); + } + if (!m41->engram_layers.empty()) { + const size_t n_cols = m41->engram_n_cols(); + int32_t pad_id = 0; // any raw id whose compressed id is the pad's + for (size_t i = 0; i < m41->engram_token_map.size(); ++i) { + if (m41->engram_token_map[i] == m41->engram_pad_cid) { + pad_id = (int32_t) i; + break; + } + } + std::vector primes(m41->engram_primes.begin(), m41->engram_primes.end()); + GGML_ASSERT(primes.size() == n_cols*m41->engram_layers.size()); + add_kv(LLM_KV_ENGRAM_LAYERS, std::vector(m41->engram_layers.begin(), m41->engram_layers.end())); + add_kv(LLM_KV_ENGRAM_NUM_EMBEDDINGS, m41->engram_n_rows); + add_kv(LLM_KV_ENGRAM_MAX_NGRAM_SIZE, m41->engram_max_ngram); + add_kv(LLM_KV_ENGRAM_HEAD_COUNT, m41->engram_n_head); + add_kv(LLM_KV_ENGRAM_HEAD_DIM, m41->engram_head_dim); + add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, (uint32_t) pad_id); + add_kv(LLM_KV_ENGRAM_TOKEN_MAP, m41->engram_token_map); + add_kv(LLM_KV_ENGRAM_MULTIPLIERS, m41->engram_mult); + add_kv(LLM_KV_ENGRAM_PRIMES, primes); + } + } + // the PLE group only means anything whole: write all of it or none if (hparams.ple_n_heads > 0) { std::vector ple_layers; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 803a8329e9db..2442a306140a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -208,6 +208,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_dots3note(params); case LLM_ARCH_DEEPSEEK4: return new llama_model_deepseek4(params); + case LLM_ARCH_DEEPSEEK41: + return new llama_model_deepseek41(params); case LLM_ARCH_GLM_DSA: return new llama_model_glm_dsa(params); case LLM_ARCH_MISTRAL4: @@ -394,7 +396,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const llama_meta_device_get_split_state_userdata * ud = (const llama_meta_device_get_split_state_userdata *) userdata; const llama_hparams & hparams = ud->model->hparams; const std::string tensor_name = tensor->name; - const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 || + const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 || ud->model->arch == LLM_ARCH_DEEPSEEK41 || (ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0); static const std::regex pattern_q_weight ("blk\\.\\d*\\.attn_q.weight"); @@ -3010,6 +3012,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr); } } break; + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_DEEPSEEK4: { GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); @@ -3419,7 +3422,7 @@ int32_t llama_model_n_head_kv(const llama_model * model) { int32_t llama_model_n_swa(const llama_model * model) { // dsv4 kv-cache has SWA but it cannot be used as a rollback because of // other compression ratios, so we return 0 here - if (model->arch == LLM_ARCH_DEEPSEEK4) { + if (model->arch == LLM_ARCH_DEEPSEEK4 || model->arch == LLM_ARCH_DEEPSEEK41) { return 0; } return model->hparams.n_swa; @@ -3506,6 +3509,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_DEEPSEEK2OCR: case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_MUSE_SPARK: case LLM_ARCH_PLM: diff --git a/src/llama-model.h b/src/llama-model.h index a187c274932b..1c4b93a09f1f 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -605,6 +605,14 @@ struct llama_layer { struct ggml_tensor * indexer_kpool_gate = nullptr; struct ggml_tensor * indexer_kpool_ape = nullptr; + // deepseek41 + struct ggml_tensor * engram_embd = nullptr; // [head_dim, n_rows], host-resident gather table + struct ggml_tensor * engram_kv = nullptr; // [n_hash_cols*head_dim, n_embd*(hc+1)] + struct ggml_tensor * engram_gate = nullptr; // [n_embd, hc], q_weight*k_weight folded + struct ggml_tensor * dspark_main_proj = nullptr; + struct ggml_tensor * dspark_main_norm = nullptr; + struct ggml_tensor * dspark_norm = nullptr; + // MSA struct ggml_tensor * index_q_proj = nullptr; struct ggml_tensor * index_k_proj = nullptr; diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp new file mode 100644 index 000000000000..78d553b7d794 --- /dev/null +++ b/src/models/deepseek41.cpp @@ -0,0 +1,1050 @@ +#include "llama-hparams.h" +#include "models.h" + +#include "llama-kv-cache-dsv4.h" + +#include +#include +#include +#include + +// DeepSeek-V4.1 (ref: deepseek-ai/DeepSeek-V4.1-Flash-1017 inference/model.py, engram.py) +// +// Differences to DeepSeek-V4 handled here: +// - Engram on a few layers: hashed 2..N-gram rows gathered from host-resident tables, mixed into every HC copy +// - single-pass mHC: a sublayer collapses the HC copies with the pre-mix computed by the PREVIOUS sublayer +// (layer 0 attention: copy 0), and there is no hc_head (the output collapses with the last FFN pre-mix) +// - compress ratios 1 and 2 (plain softmax pooling, no ape / overlap), compressed KV and index K produced only by +// the kv source layers and read by the layers after them; index K = k_norm(wk(latent)), stored next to the KV row +// - indexer top-k computed only on index source layers, reused by the layers in between +// - two-level candidate selection: the candidate source layer picks blocks, later index layers score inside them +// - no per-head q norm after wq_b, no hash-routed (tid2eid) layers + +static float dsv41_rope_attn_factor(float freq_scale, float ext_factor) { + if (ext_factor == 0.0f) { + return 1.0f; + } + + return 1.0f / (1.0f + 0.1f*logf(1.0f/freq_scale)); +} + +static constexpr int64_t DSV41_ENGRAM_DEAD = -1; + +// +// loading +// + +void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { + 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_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); + 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_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_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + + 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_ATTENTION_COMPRESS_ROPE_FREQ_BASE, hparams.dsv4_compress_rope_base); + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + 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); + hparams.dsv4_hash_layer_count = 0; + + 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_all) { + throw std::runtime_error("DeepSeek-V4.1 compress_ratios is shorter than block_count"); + } + GGML_ASSERT(n_compress_ratios <= LLAMA_MAX_LAYERS); + ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios); + + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { + throw std::runtime_error("DeepSeek-V4.1 loader expects sqrtsoftplus MoE scoring"); + } + + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + hparams.set_swa_pattern(0); + hparams.non_causal_type = LLAMA_NON_CAUSAL_TYPE_SWA_FULL; + for (uint32_t il = hparams.n_layer(); il < hparams.n_layer_all; ++il) { + hparams.is_swa_impl[il] = true; + } + + // cross-layer sharing: every compressed layer reads the latest kv source at or before it, and attends with the + // top-k of the latest index source at or before it (SharedAttentionRuntime in the reference) + std::vector kv_src; + std::vector idx_src; + ml.get_arr(LLM_KV_ATTENTION_KV_SOURCE_LAYERS, kv_src); + ml.get_arr(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, idx_src); + + const uint32_t n_layer = hparams.n_layer(); + int32_t last_kv = -1; + int32_t last_idx = -1; + for (uint32_t il = 0; il < LLAMA_MAX_LAYERS; ++il) { + hparams.dsv41_kv_src[il] = -1; + hparams.dsv41_topk_src[il] = -1; + } + for (uint32_t il = 0; il < n_layer; ++il) { + const uint32_t ratio = hparams.dsv4_compress_ratios[il]; + if (std::find(kv_src.begin(), kv_src.end(), il) != kv_src.end()) { + if (ratio == 0) { + throw std::runtime_error(format("DeepSeek-V4.1 kv source layer %u has compress ratio 0", il)); + } + last_kv = (int32_t) il; + } + if (std::find(idx_src.begin(), idx_src.end(), il) != idx_src.end()) { + if (ratio == 0) { + throw std::runtime_error(format("DeepSeek-V4.1 index source layer %u has compress ratio 0", il)); + } + last_idx = (int32_t) il; + } + if (ratio == 0) { + continue; + } + if (ratio != 1 && ratio != 2) { + throw std::runtime_error(format("DeepSeek-V4.1 layer %u: compress ratio %u (only 0, 1, 2 supported)", il, ratio)); + } + if (last_kv < 0 || last_idx < 0 || + hparams.dsv4_compress_ratios[last_kv] != ratio || hparams.dsv4_compress_ratios[last_idx] != ratio || + last_idx < last_kv) { + throw std::runtime_error(format("DeepSeek-V4.1 layer %u: no matching kv/index source before it", il)); + } + hparams.dsv41_kv_src[il] = last_kv; + hparams.dsv41_topk_src[il] = last_idx; + } + + hparams.dsv41_cand_src = -1; + hparams.dsv41_cand_topk_blocks = 0; + hparams.dsv41_cand_block_size = 0; + ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, hparams.dsv41_cand_src, false); + if (hparams.dsv41_cand_src >= 0) { + ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, hparams.dsv41_cand_topk_blocks); + ml.get_key(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, hparams.dsv41_cand_block_size); + const int32_t cs = hparams.dsv41_cand_src; + if ((uint32_t) cs >= n_layer || hparams.dsv41_topk_src[cs] != cs || hparams.dsv41_cand_block_size == 0 || + hparams.dsv41_cand_topk_blocks == 0) { + throw std::runtime_error("DeepSeek-V4.1 candidate source layer must be an index source"); + } + for (uint32_t il = cs + 1; il < n_layer; ++il) { + if (hparams.dsv41_topk_src[il] == (int32_t) il && hparams.dsv4_compress_ratios[il] != hparams.dsv4_compress_ratios[cs]) { + throw std::runtime_error("DeepSeek-V4.1 candidate consumers must share the source's compress ratio"); + } + } + } + + // Engram + std::vector e_layers; + ml.get_arr(LLM_KV_ENGRAM_LAYERS, e_layers, false); + if (!e_layers.empty()) { + uint32_t pad_id = 0; + std::vector primes; + ml.get_arr(LLM_KV_ENGRAM_NUM_EMBEDDINGS, engram_n_rows); + ml.get_key(LLM_KV_ENGRAM_MAX_NGRAM_SIZE, engram_max_ngram); + ml.get_key(LLM_KV_ENGRAM_HEAD_COUNT, engram_n_head); + ml.get_key(LLM_KV_ENGRAM_HEAD_DIM, engram_head_dim); + ml.get_key(LLM_KV_ENGRAM_PAD_TOKEN_ID, pad_id); + ml.get_arr(LLM_KV_ENGRAM_TOKEN_MAP, engram_token_map); + ml.get_arr(LLM_KV_ENGRAM_MULTIPLIERS, engram_mult); + ml.get_arr(LLM_KV_ENGRAM_PRIMES, primes); + + const size_t n_e = e_layers.size(); + const size_t n_cols = (size_t) (engram_max_ngram - 1)*engram_n_head; + if (engram_max_ngram < 2 || engram_n_head == 0 || engram_n_rows.size() != n_e || + engram_mult.size() != n_e*engram_max_ngram || primes.size() != n_e*n_cols || + pad_id >= engram_token_map.size()) { + throw std::runtime_error("DeepSeek-V4.1 engram metadata is inconsistent"); + } + engram_layers.assign(e_layers.begin(), e_layers.end()); + hparams.dsv41_tok_hist = true; // state save/restore must carry the n-gram history + hparams.dsv41_hist_len = engram_max_ngram > 0 ? (uint32_t) engram_max_ngram - 1 : 0; + engram_pad_cid = engram_token_map[pad_id]; + engram_primes.assign(primes.begin(), primes.end()); + engram_offsets.resize(primes.size()); + for (size_t e = 0; e < n_e; ++e) { + uint64_t off = 0; + for (size_t c = 0; c < n_cols; ++c) { + engram_offsets[e*n_cols + c] = off; + off += engram_primes[e*n_cols + c]; + } + if (off > engram_n_rows[e] || off > (uint64_t) INT32_MAX) { + throw std::runtime_error(format("DeepSeek-V4.1 engram layer %u: buckets (%llu) exceed the table rows (%u)", + e_layers[e], (unsigned long long) off, engram_n_rows[e])); + } + } + for (uint32_t il : e_layers) { + if (il >= n_layer) { + throw std::runtime_error(format("DeepSeek-V4.1 engram layer %u is out of range", il)); + } + } + } + + type = LLM_TYPE_UNKNOWN; +} + +void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + 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; + const int64_t n_idx_embd = hparams.indexer_head_size; + const int64_t n_idx_head = hparams.indexer_n_head; + + // DSpark stages (blk.n_layer..) are loaded as optional tensors and unused here: the draft runs as a DFLASH model + uint32_t n_expert_dspark = n_expert; + ml.get_key(LLM_KV_DSPARK_EXPERT_COUNT, n_expert_dspark, false); + + 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); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + for (int i = 0; i < n_layer_all; ++i) { + auto & layer = layers[i]; + const bool mtp = i >= n_layer; + const int flags = mtp ? TENSOR_NOT_REQUIRED : 0; + const int64_t n_exp = mtp ? (int64_t) n_expert_dspark : n_expert; + + 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 | TENSOR_ALLOW_RESHAPE); + 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); + + if (!mtp && hparams.dsv41_kv_src[i] == i) { + const uint32_t ratio = hparams.dsv4_compress_ratios[i]; + layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, n_embd_head}, 0); + layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, 0); + if (ratio > 1) { + layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, n_embd_head}, 0); + } + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd_head, n_idx_embd}, 0); + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_idx_embd}, 0); + } + if (!mtp && hparams.dsv41_topk_src[i] == i) { + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_idx_head * n_idx_embd}, 0); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_idx_head}, 0); + } + + if (!mtp) { + const int e = engram_index(i); + if (e >= 0) { + const int64_t n_cols = engram_n_cols(); + layer.engram_embd = create_tensor(tn(LLM_TENSOR_ENGRAM_EMBD, "weight", i), {(int64_t) engram_head_dim, (int64_t) engram_n_rows[e]}, 0); + layer.engram_kv = create_tensor(tn(LLM_TENSOR_ENGRAM_KV, "weight", i), {n_cols * engram_head_dim, n_embd * (hc_mult + 1)}, 0); + layer.engram_gate = create_tensor(tn(LLM_TENSOR_ENGRAM_GATE, "weight", i), {n_embd, hc_mult}, 0); + } + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_exp}, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_exp}, flags); + layer.ffn_exp_probs_b_vl = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B_VL, "bias", i), {n_exp}, TENSOR_NOT_REQUIRED); + 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_exp}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_exp}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_exp}, 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 (mtp) { + layer.dspark_main_proj = create_tensor(tn(LLM_TENSOR_DSPARK_MAIN_PROJ, "weight", i), {n_embd * 3, n_embd}, TENSOR_NOT_REQUIRED); + layer.dspark_main_norm = create_tensor(tn(LLM_TENSOR_DSPARK_MAIN_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); + layer.dspark_norm = create_tensor(tn(LLM_TENSOR_DSPARK_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); + } + } +} + +std::unique_ptr llama_model_deepseek41::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + throw std::runtime_error("DeepSeek-V4.1 has no MTP graph: its DSpark draft is a separate DFLASH GGUF " + "(convert_hf_to_gguf.py --dspark), used with --spec-type draft-dspark"); + } + return std::make_unique(*this, params); +} + +// +// Engram n-gram hash (engram.py NgramHashState.forward) +// + +void llama_model_deepseek41::engram_hash(int e, const llama_ubatch & ubatch, const std::vector & prev, std::vector & out) const { + const int64_t n_tokens = ubatch.n_tokens; + const int64_t n_gram = engram_max_ngram; + const int64_t n_prev = n_gram - 1; + const int64_t n_head = engram_n_head; + const int64_t n_cols = engram_n_cols(); + const int64_t n_map = (int64_t) engram_token_map.size(); + + const uint64_t * mult = engram_mult.data() + (size_t) e*n_gram; + const uint64_t * primes = engram_primes.data() + (size_t) e*n_cols; + const uint64_t * offs = engram_offsets.data() + (size_t) e*n_cols; + + out.resize(n_tokens*n_cols); + + // compressed id of a raw token; embeddings (image spans) and unknown history are DEAD + const auto cid = [&](llama_token t) -> int64_t { + if (t < 0 || t >= n_map) { + return DSV41_ENGRAM_DEAD; + } + return engram_token_map[t]; + }; + + std::vector toks(n_gram); + for (int64_t i = 0; i < n_tokens; ++i) { + const bool is_embd = !ubatch.token || (ubatch.is_mixed() && ubatch.type[i]); + const llama_pos pos = ubatch.pos[i]; + + // shift 0 is the token itself; once a shift is blocked (sequence start or dead token) every longer one is too + int64_t c = is_embd ? DSV41_ENGRAM_DEAD : cid(ubatch.token[i]); + bool blocked = c == DSV41_ENGRAM_DEAD; + toks[0] = blocked ? engram_pad_cid : c; + for (int64_t s = 1; s < n_gram; ++s) { + const llama_token t = prev[i*n_prev + (n_prev - s)]; + c = (pos - s < 0) ? DSV41_ENGRAM_DEAD : cid(t); + blocked = blocked || c == DSV41_ENGRAM_DEAD; + toks[s] = blocked ? engram_pad_cid : c; + } + + uint64_t rolling = toks[0]*mult[0]; + for (int64_t o = 1; o < n_gram; ++o) { + rolling ^= toks[o]*mult[o]; + for (int64_t h = 0; h < n_head; ++h) { + const int64_t col = (o - 1)*n_head + h; + out[i*n_cols + col] = (int32_t) (rolling % primes[col] + offs[col]); + } + } + } +} + +// Host-side inputs: Engram row ids per engram layer, and the candidate-block pin (+1e30 on the block holding each +// query's newest visible compressed row: select_candidate_blocks pins it so a partly filled block always survives). +class llm_graph_input_dsv41 : public llm_graph_input_i { +public: + llm_graph_input_dsv41(const llama_model_deepseek41 & model, const llama_kv_cache_dsv4_context * mctx) : + model(model), mctx(mctx) {} + virtual ~llm_graph_input_dsv41() = default; + + void set_input(const llama_ubatch * ubatch) override; + bool can_reuse(const llm_graph_params & params) override; + + std::vector engram_ids; // per engram layer: I32 [n_cols*n_tokens] + ggml_tensor * engram_keep = nullptr; // F32 [1, 1, n_tokens], 0 for embedding (image) tokens; only in embd ubatches + ggml_tensor * cand_pin = nullptr; // F32 [n_blk, n_tokens/n_stream, 1, n_stream] + bool cand_csa = false; // candidate source reads the CSA (ratio 2) slot, else HCA (ratio 1) + + const llama_model_deepseek41 & model; + const llama_kv_cache_dsv4_context * mctx; + + std::vector prev; + std::vector ids; + std::vector pin; + std::vector keep; + +private: + const llama_kv_cache_dsv4_context::comp_plan & cand_plan(const llama_ubatch & ubatch) const { + return cand_csa ? mctx->get_csa_plan(ubatch) : mctx->get_hca_plan(ubatch); + } +}; + +void llm_graph_input_dsv41::set_input(const llama_ubatch * ubatch) { + if (!engram_ids.empty()) { + // a token shared by several sequences (n_seq_id > 1) reads the history of its first sequence: + // llama_kv_cache_dsv4::init_batch rejected the batch unless every sequence of the token has the same + // predecessors, so the first one is representative + // predecessors come from the raw window cells, which already hold this ubatch + mctx->get_raw()->get_prev_tokens(*ubatch, model.engram_max_ngram - 1, prev); + for (size_t e = 0; e < engram_ids.size(); ++e) { + if (!engram_ids[e]->buffer) { + continue; + } + model.engram_hash((int) e, *ubatch, prev, ids); + ggml_backend_tensor_set(engram_ids[e], ids.data(), 0, ids.size()*sizeof(int32_t)); + } + } + + // image tokens take no part in an n-gram (hashed DEAD above) and get no Engram contribution (gate 0), as model.py + if (engram_keep && engram_keep->buffer) { + keep.resize(ubatch->n_tokens); + for (uint32_t i = 0; i < ubatch->n_tokens; ++i) { + const bool is_embd = !ubatch->token || (ubatch->is_mixed() && ubatch->type[i]); + keep[i] = is_embd ? 0.0f : 1.0f; + } + ggml_backend_tensor_set(engram_keep, keep.data(), 0, keep.size()*sizeof(float)); + } + + if (cand_pin && cand_pin->buffer) { + const auto & plan = cand_plan(*ubatch); + const int64_t n_blk = cand_pin->ne[0]; + const int64_t bs = model.hparams.dsv41_cand_block_size; + GGML_ASSERT((int64_t) plan.n_visible.size() == (int64_t) ubatch->n_tokens); + pin.assign(n_blk*ubatch->n_tokens, 0.0f); + for (int64_t i = 0; i < (int64_t) ubatch->n_tokens; ++i) { + const int64_t nv = plan.n_visible[i]; + if (nv > 0) { + const int64_t b = (nv - 1)/bs; + GGML_ASSERT(b < n_blk); + pin[i*n_blk + b] = 1e30f; + } + } + ggml_backend_tensor_set(cand_pin, pin.data(), 0, pin.size()*sizeof(float)); + } +} + +bool llm_graph_input_dsv41::can_reuse(const llm_graph_params & params) { + mctx = static_cast(params.mctx); + + bool res = true; + for (auto * t : engram_ids) { + res &= t->ne[0] == (int64_t) model.engram_n_cols()*params.ubatch.n_tokens; + } + res &= (engram_keep != nullptr) == (!engram_ids.empty() && params.ubatch.embd != nullptr); + if (engram_keep) { + res &= engram_keep->ne[2] == (int64_t) params.ubatch.n_tokens; + } + if (cand_pin) { + const auto & plan = cand_plan(params.ubatch); + const int64_t bs = model.hparams.dsv41_cand_block_size; + res &= plan.n_kv > 0 && cand_pin->ne[0] == (plan.n_kv + bs - 1)/bs; + res &= cand_pin->ne[1]*cand_pin->ne[3] == (int64_t) params.ubatch.n_tokens; + } + return res; +} + +// +// graph helpers +// + +static ggml_tensor * dsv41_view_2d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t i0) { + return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], ggml_row_size(t->type, i0)); +} + +static ggml_tensor * dsv41_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) { + return ggml_view_1d(ctx, t, ne0, ggml_row_size(t->type, i0)); +} + +ggml_tensor * llama_model_deepseek41::graph::build_hc_mixes( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) const { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc*n_embd; + const int64_t hc_mix_dim = (2 + hc)*hc; + const int64_t nt = x->ne[2]; + + GGML_ASSERT(hc == 4); + GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim); + + // normalized over the whole flattened hc*d stream, one statistic per token + ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); + ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps); + ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); + // F32 weights (model.py hc_pre runs in fp32): no TF32 rounding on CUDA, else the mixes depend on the ubatch size. + // Skipped for quantized weights: third-party GGUFs quantize hc_*_fn, and a quantized src0 with PREC_F32 asserts in MMQ. + if (!ggml_is_quantized(hc_fn->type)) { + ggml_prec_set_src(mixes, GGML_PREC_F32, 1); + } + cb(mixes, "hc_mixes", il); + + ggml_tensor * pre = dsv41_view_2d(ctx0, mixes, hc, nt, 0); + pre = ggml_mul(ctx0, pre, dsv41_view_1d(ctx0, hc_scale, 1, 0)); + pre = ggml_add(ctx0, pre, dsv41_view_1d(ctx0, hc_base, hc, 0)); + pre = ggml_sigmoid(ctx0, pre); + pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); + cb(pre, "hc_pre", il); + + *post = dsv41_view_2d(ctx0, mixes, hc, nt, hc); + *post = ggml_mul(ctx0, *post, dsv41_view_1d(ctx0, hc_scale, 1, 1)); + *post = ggml_add(ctx0, *post, dsv41_view_1d(ctx0, hc_base, hc, hc)); + *post = ggml_sigmoid(ctx0, *post); + *post = ggml_scale(ctx0, *post, 2.0f); + cb(*post, "hc_post", il); + + if (cparams.fused_dsv4_hc_comb) { + *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps, + (int32_t) hparams.dsv4_hc_sinkhorn_iters); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il}); + } else { + *comb = dsv41_view_2d(ctx0, mixes, hc*hc, nt, 2*hc); + *comb = ggml_mul(ctx0, *comb, dsv41_view_1d(ctx0, hc_scale, 1, 2)); + *comb = ggml_add(ctx0, *comb, dsv41_view_1d(ctx0, hc_base, hc*hc, 2*hc)); + *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt); + *comb = build_hc_sinkhorn(*comb, il); + } + cb(*comb, "hc_comb", il); + + return pre; +} + +// Engram.forward: x [n_embd, hc, n_tokens], ids I32 [n_cols*n_tokens] +ggml_tensor * llama_model_deepseek41::graph::build_engram( + const llama_model_deepseek41 & model, + ggml_tensor * x, + ggml_tensor * ids, + ggml_tensor * keep, + int il) const { + const auto & layer = model.layers[il]; + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[2]; + const int64_t n_cols = model.engram_n_cols(); + const int64_t hd = model.engram_head_dim; + + ggml_tensor * rows = ggml_get_rows(ctx0, layer.engram_embd, ids); // [hd, n_cols*nt] + if (rows->type != GGML_TYPE_F32) { + rows = ggml_cast(ctx0, rows, GGML_TYPE_F32); + } + rows = ggml_reshape_2d(ctx0, rows, hd*n_cols, nt); + cb(rows, "engram_rows", il); + + ggml_tensor * kv = build_lora_mm(layer.engram_kv, rows); // [n_embd*(hc+1), nt] + cb(kv, "engram_kv", il); + + ggml_tensor * key = ggml_view_3d(ctx0, kv, n_embd, hc, nt, ggml_row_size(kv->type, n_embd), kv->nb[1], 0); + ggml_tensor * value = ggml_view_3d(ctx0, kv, n_embd, 1, nt, ggml_row_size(kv->type, n_embd), kv->nb[1], ggml_row_size(kv->type, hc*n_embd)); + key = ggml_cont(ctx0, key); + value = ggml_cont(ctx0, value); + + // dot_c = sum_d (x * rstd_x) * (key * rstd_k) * w / sqrt(n_embd), normalized per (token, copy) + ggml_tensor * xn = ggml_rms_norm(ctx0, x, norm_rms_eps); + ggml_tensor * kn = ggml_rms_norm(ctx0, key, norm_rms_eps); + ggml_tensor * dot = ggml_mul(ctx0, ggml_mul(ctx0, xn, kn), layer.engram_gate); + dot = ggml_sum_rows(ctx0, dot); // [1, hc, nt] + dot = ggml_scale(ctx0, dot, 1.0f/sqrtf((float) n_embd)); + cb(dot, "engram_dot", il); + + // gate = sigmoid(sign(dot) * sqrt(max(|dot|, 1e-6))) + ggml_tensor * mag = ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, dot), 1e-6f, INFINITY)); + ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_mul(ctx0, mag, ggml_sgn(ctx0, dot))); + if (keep) { + gate = ggml_mul(ctx0, gate, keep); // [1, hc, nt] * [1, 1, nt]: image tokens pass through untouched + } + cb(gate, "engram_gate", il); + + ggml_tensor * upd = ggml_mul(ctx0, ggml_repeat_4d(ctx0, value, n_embd, hc, nt, 1), gate); + return ggml_add(ctx0, x, upd); +} + +static void dsv41_state_restore( + ggml_context * ctx, + const llm_graph_input_dsv4::comp_input & inp, + const llama_dsv4_comp_state * state, + int32_t il, + ggml_tensor ** kv, + ggml_tensor ** score) { + *kv = state->get_kv_all(ctx, il); + *score = state->get_score_all(ctx, il); + + if (inp.state_restore_src_idxs == nullptr || inp.state_restore_dst_idxs == nullptr) { + return; + } + + *kv = state->cpy_kv (ctx, ggml_get_rows(ctx, *kv, inp.state_restore_src_idxs), inp.state_restore_dst_idxs, il); + *score = state->cpy_score(ctx, ggml_get_rows(ctx, *score, inp.state_restore_src_idxs), inp.state_restore_dst_idxs, il); +} + +ggml_tensor * llama_model_deepseek41::graph::build_attention41( + const llama_model_deepseek41 & model, + llm_graph_input_dsv4 * inp_dsv4, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * cand_pin, + int il) { + const auto & layer = model.layers[il]; + llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw(); + + 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; + const int64_t n_groups = hparams.dsv4_o_group_count; + const int64_t n_heads_group = n_head / n_groups; + const int64_t o_lora_rank = hparams.dsv4_o_lora_rank; + const int64_t o_group_dim = n_heads_group*n_embd_head; + const int64_t n_idx_embd = hparams.indexer_head_size; + const int64_t n_idx_head = hparams.indexer_n_head; + const int64_t nt = cur->ne[1]; + const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + + GGML_ASSERT(n_embd_head == n_embd_head_v); + GGML_ASSERT(n_head % n_groups == 0); + + const int64_t ratio = hparams.dsv4_compress_ratios[il]; + + // layers with compressed KV use the compress theta + YaRN, window-only layers plain rope + const bool use_compress_rope = ratio != 0; + const float freq_base_l = use_compress_rope ? hparams.dsv4_compress_rope_base : freq_base; + const float freq_scale_l = use_compress_rope ? freq_scale : 1.0f; + const float ext_factor_l = use_compress_rope ? ext_factor : 0.0f; + const float attn_factor_l = dsv41_rope_attn_factor(freq_scale_l, ext_factor_l); + const float beta_fast_l = use_compress_rope ? beta_fast : 0.0f; + const float beta_slow_l = use_compress_rope ? beta_slow : 0.0f; + const int32_t n_ctx_orig_l = use_compress_rope ? n_ctx_orig : 0; + + const auto rope = [&](ggml_tensor * t, ggml_tensor * pos, int64_t n_nope) { + t = ggml_rope_ext(ctx0, t, pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, + freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); + return ggml_rope_set_offset(t, n_nope); + }; + + ggml_tensor * qr = build_lora_mm(layer.wq_a, cur); + qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(qr, "qr_norm", il); + + ggml_tensor * q = build_lora_mm(layer.wq_b, qr); + q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, nt); + q = rope(q, inp_pos, n_embd_head_nope); + cb(q, "q", il); + + ggml_tensor * kv = build_lora_mm(layer.wkv, cur); + kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il); + kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, nt); + kv = rope(kv, inp_pos, n_embd_head_nope); + cb(kv, "kv", il); + + ggml_tensor * out = nullptr; + + if (ratio == 0) { + out = build_raw_attention(inp_attn, q, kv, layer.attn_sinks, kq_scale, il); + } else { + GGML_ASSERT(ratio == 1 || ratio == 2); + const bool use_csa = ratio == 2; + const auto & ci = use_csa ? inp_dsv4->get_csa() : inp_dsv4->get_hca(); + const auto * comp_ctx = use_csa ? inp_dsv4->mctx->get_csa() : inp_dsv4->mctx->get_hca(); + const auto * comp_state = use_csa ? inp_dsv4->mctx->get_csa_state() : inp_dsv4->mctx->get_hca_state(); + const int32_t src = hparams.dsv41_kv_src[il]; + GGML_ASSERT(src >= 0 && (int64_t) comp_state->get_ratio() == ratio); + + // ---- compressor (kv sources only): pool ratio tokens, norm, index K, rope, write [kv | index K] + if (src == il && ci.state_pos) { + ggml_tensor * state_kv = build_lora_mm(layer.attn_comp_wkv, cur); + cb(state_kv, "comp_state_kv", il); + ggml_tensor * state_score = ratio > 1 ? build_lora_mm(layer.attn_comp_wgate, cur) : ggml_scale(ctx0, state_kv, 0.0f); + cb(state_score, "comp_state_score", il); + + ggml_tensor * restored_kv = nullptr; + ggml_tensor * restored_score = nullptr; + dsv41_state_restore(ctx0, ci, comp_state, il, &restored_kv, &restored_score); + + if (ci.state_write_idxs) { + const int64_t n_blocks = ci.state_write_pos->ne[0]; + GGML_ASSERT(ci.state_read_idxs->ne[0] == ratio*n_blocks); + + ggml_tensor * src_kv = ggml_concat(ctx0, + dsv41_view_2d(ctx0, restored_kv, restored_kv->ne[0], comp_state->get_n_rows(), 0), state_kv, 1); + ggml_tensor * src_score = ggml_concat(ctx0, + dsv41_view_2d(ctx0, restored_score, restored_score->ne[0], comp_state->get_n_rows(), 0), state_score, 1); + + ggml_tensor * v = ggml_get_rows(ctx0, src_kv, ci.state_read_idxs); + ggml_tensor * s = ggml_get_rows(ctx0, src_score, ci.state_read_idxs); + v = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, v, n_embd_head, ratio, n_blocks), 1, 0, 2, 3)); + s = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, s, n_embd_head, ratio, n_blocks), 1, 0, 2, 3)); + + // per-channel softmax over the ratio slots of a group + ggml_tensor * comp = ggml_sum_rows(ctx0, ggml_mul(ctx0, v, ggml_soft_max(ctx0, s))); // [1, hd, nb] + comp = ggml_cont(ctx0, ggml_permute(ctx0, comp, 1, 0, 2, 3)); // [hd, 1, nb] + comp = build_norm(comp, layer.attn_comp_norm, nullptr, LLM_NORM_RMS, il); + cb(comp, "comp_latent", il); + + // index K from the RoPE-free latent; a latent stands for the first token of its group + ggml_tensor * ik = build_lora_mm(layer.indexer_attn_k, comp); // [n_idx_embd, 1, nb] + ik = build_norm(ik, layer.indexer_k_norm, nullptr, LLM_NORM_RMS, il); + ik = rope(ik, ci.state_write_pos, n_idx_embd - n_embd_head_rope); + cb(ik, "comp_index_k", il); + + comp = rope(comp, ci.state_write_pos, n_embd_head_nope); + cb(comp, "comp_kv", il); + + ggml_tensor * row = ggml_concat(ctx0, comp, ik, 0); // [hd + n_idx_embd, 1, nb] + ggml_build_forward_expand(gf, comp_ctx->cpy_k(ctx0, row, ci.state_write_idxs, il)); + } + + // rollback snapshots and the persistent partial-group state (as V4 HCA) + if (ci.state_snapshot_src_idxs && ci.state_snapshot_dst_idxs) { + ggml_tensor * snap_kv = ggml_concat(ctx0, restored_kv, state_kv, 1); + ggml_tensor * snap_score = ggml_concat(ctx0, restored_score, state_score, 1); + ggml_build_forward_expand(gf, comp_state->cpy_kv(ctx0, + ggml_get_rows(ctx0, snap_kv, ci.state_snapshot_src_idxs), ci.state_snapshot_dst_idxs, il)); + ggml_build_forward_expand(gf, comp_state->cpy_score(ctx0, + ggml_get_rows(ctx0, snap_score, ci.state_snapshot_src_idxs), ci.state_snapshot_dst_idxs, il)); + } + + ggml_build_forward_expand(gf, comp_state->cpy_kv(ctx0, + ggml_get_rows(ctx0, state_kv, ci.state_persist_src_idxs), ci.state_persist_dst_idxs, il)); + ggml_build_forward_expand(gf, comp_state->cpy_score(ctx0, + ggml_get_rows(ctx0, state_score, ci.state_persist_src_idxs), ci.state_persist_dst_idxs, il)); + } + + if (!ci.kq_mask) { + // no compressed row visible to any token of this ubatch: window only + out = build_raw_attention(inp_attn, q, kv, layer.attn_sinks, kq_scale, il); + } else { + const int64_t n_comp = ci.kq_mask->ne[0]; + + ggml_tensor * comp_all = comp_ctx->get_k(ctx0, src); // [hd + n_idx_embd, 1, n_cells, n_stream] + GGML_ASSERT(comp_all->ne[0] == n_embd_head + n_idx_embd); + GGML_ASSERT(n_comp <= comp_all->ne[2]); + const int64_t n_stream = comp_all->ne[3]; + + ggml_tensor * comp_mask = ci.kq_mask; + if (comp_mask->type != GGML_TYPE_F32) { + comp_mask = ggml_cast(ctx0, comp_mask, GGML_TYPE_F32); + } + + // ---- top-k: own indexer on index sources, reused otherwise + const int32_t topk_src = hparams.dsv41_topk_src[il]; + ggml_tensor * top_k = nullptr; + if (topk_src == il) { + ggml_tensor * ik = ggml_view_4d(ctx0, comp_all, n_idx_embd, 1, n_comp, n_stream, + comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], ggml_row_size(comp_all->type, n_embd_head)); + ik = ggml_cont(ctx0, ik); + cb(ik, "idx_k", il); + + ggml_tensor * iq = build_lora_mm(layer.indexer_attn_q_b, qr); + iq = ggml_reshape_3d(ctx0, iq, n_idx_embd, n_idx_head, nt); + iq = rope(iq, inp_pos, n_idx_embd - n_embd_head_rope); + cb(iq, "idx_q", il); + + ggml_tensor * iw = build_lora_mm(layer.indexer_proj, cur); + iw = ggml_scale(ctx0, iw, 1.0f/sqrtf(float(n_idx_embd*n_idx_head))); + cb(iw, "idx_w", il); + + iq = ggml_view_4d(ctx0, iq, iq->ne[0], iq->ne[1], nt/n_stream, n_stream, + iq->nb[1], iq->nb[2], iq->nb[2]*(nt/n_stream), 0); + iw = ggml_view_4d(ctx0, iw, iw->ne[0], nt/n_stream, 1, n_stream, + iw->nb[1], iw->nb[1]*(nt/n_stream), iw->nb[1]*(nt/n_stream), 0); + + // score[j, t] = sum_h relu(q_h . k_j) * w_h + ggml_tensor * kq = ggml_mul_mat(ctx0, + ggml_permute(ctx0, ik, 0, 2, 1, 3), // [d, n_comp, 1, ns] + ggml_permute(ctx0, iq, 0, 2, 1, 3)); // [d, nt/ns, n_ih, ns] -> [n_comp, nt/ns, n_ih, ns] + if (!ggml_is_quantized(ik->type)) { + ggml_prec_set_src(kq, GGML_PREC_F32, 1); // F32 index keys: no TF32 rounding of the near-tie scores + } + kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); // [n_ih, nt/ns, n_comp, ns] + ggml_tensor * score = ggml_sum_rows(ctx0, ggml_mul(ctx0, ggml_relu(ctx0, kq), iw)); // [1, nt/ns, n_comp, ns] + score = ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3)); // [n_comp, nt/ns, 1, ns] + score = ggml_add(ctx0, score, comp_mask); + cb(score, "idx_score", il); + + const int32_t cand_src = hparams.dsv41_cand_src; + if (cand_src == il && cand_pin) { + // level one: best position per block, pin the newest block, keep the top blocks + const int64_t bs = hparams.dsv41_cand_block_size; + const int64_t n_blk = (n_comp + bs - 1)/bs; + GGML_ASSERT(cand_pin->ne[0] == n_blk); + + ggml_tensor * s2 = ggml_reshape_2d(ctx0, score, n_comp, nt); + if (n_blk*bs != n_comp) { + ggml_tensor * pad = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_blk*bs - n_comp, nt); + s2 = ggml_concat(ctx0, s2, ggml_fill(ctx0, pad, -INFINITY), 0); + } + ggml_tensor * bmax = ggml_pool_1d(ctx0, s2, GGML_OP_POOL_MAX, bs, bs, 0); // [n_blk, nt] + bmax = ggml_add(ctx0, bmax, ggml_reshape_2d(ctx0, cand_pin, n_blk, nt)); + cb(bmax, "cand_block_score", il); + + const int64_t k_b = std::min(hparams.dsv41_cand_topk_blocks, n_blk); + ggml_tensor * bsel = ggml_top_k(ctx0, bmax, k_b); // [k_b, nt] + + ggml_tensor * bm = ggml_fill(ctx0, ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_blk, nt, 1), -INFINITY); + ggml_tensor * zeros = ggml_fill(ctx0, ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, k_b, nt, 1), 0.0f); + bm = ggml_set_rows(ctx0, bm, zeros, ggml_reshape_3d(ctx0, bsel, k_b, nt, 1)); // [1, n_blk, nt, 1] + bm = ggml_repeat_4d(ctx0, bm, bs, n_blk, nt, 1); + bm = ggml_reshape_2d(ctx0, bm, bs*n_blk, nt); + bm = ggml_cont(ctx0, ggml_view_2d(ctx0, bm, n_comp, nt, bm->nb[1], 0)); + cand_mask = ggml_reshape_4d(ctx0, bm, n_comp, nt/n_stream, 1, n_stream); + cb(cand_mask, "cand_mask", il); + } else if (cand_src >= 0 && cand_src < il && cand_mask) { + // level two: own scores, restricted to the source's candidate blocks + GGML_ASSERT(cand_mask->ne[0] == n_comp); + score = ggml_add(ctx0, score, cand_mask); + cb(score, "idx_score_cand", il); + } + + const int64_t k = std::min(n_comp, hparams.indexer_top_k); + top_k = ggml_cont(ctx0, ggml_top_k(ctx0, score, k)); + cb(top_k, "idx_top_k", il); + topk_of[il] = top_k; + } else { + GGML_ASSERT(topk_src >= 0); + top_k = topk_of[topk_src]; + GGML_ASSERT(top_k && "DeepSeek-V4.1: top-k source layer produced no top-k"); + } + + { + // per-layer alias of the top-k this layer attends with (reused on consumers), for selection parity + ggml_tensor * sel = ggml_view_2d(ctx0, top_k, top_k->ne[0], top_k->ne[1]*top_k->ne[2]*top_k->ne[3], top_k->nb[1], 0); + cb(sel, "sel_top_k", il); + ggml_build_forward_expand(gf, sel); + } + + // ---- sparse attention over [window raw KV | selected compressed rows] + ggml_tensor * k_rot = inp_attn->self_k_rot; + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + kv = llama_mul_mat_hadamard(ctx0, kv, k_rot); + } + + ggml_build_forward_expand(gf, q); + ggml_build_forward_expand(gf, kv); + + const llama_kv_cache_dsv4_raw_context * mctx_raw = inp_attn->mctx; + ggml_build_forward_expand(gf, mctx_raw->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il)); + + ggml_tensor * raw_k = mctx_raw->get_k(ctx0, il); + + ggml_tensor * comp_k = ggml_view_4d(ctx0, comp_all, n_embd_head, 1, n_comp, n_stream, + comp_all->nb[1], comp_all->nb[2], comp_all->nb[3], 0); + if (comp_k->type != raw_k->type) { + comp_k = ggml_cast(ctx0, comp_k, raw_k->type); + } else { + comp_k = ggml_cont(ctx0, comp_k); + } + if (k_rot) { + comp_k = llama_mul_mat_hadamard(ctx0, comp_k, k_rot); + } + cb(comp_k, "comp_k", il); + + ggml_tensor * k_all = ggml_concat(ctx0, raw_k, comp_k, 2); + + ggml_tensor * raw_mask = inp_attn->get_kq_mask(); + ggml_tensor * sel_mask = build_top_k_mask(ci.kq_mask, top_k, "comp_top_k_mask", il); + ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, sel_mask, 0); + cb(kq_mask, "comp_kq_mask", il); + + out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, layer.attn_sinks, nullptr, 0, kq_scale, il); + if (k_rot) { + out = llama_mul_mat_hadamard(ctx0, out, k_rot); + } + cb(out, "attn_comp", il); + } + } + + // undo the query rotation on the rope dims, then the grouped low-rank output projection + out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, nt); + out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, + freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); + out = ggml_rope_set_offset(out, n_embd_head_nope); + cb(out, "attn_derope", il); + + out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, nt); + out = ggml_permute(ctx0, out, 0, 2, 1, 3); + ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out); + cb(oa, "attn_wo_a", il); + oa = ggml_permute(ctx0, oa, 0, 2, 1, 3); + oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, nt); + + out = build_lora_mm(layer.wo_b, oa); + cb(out, "attn_out", il); + + return out; +} + +llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_graph_params & params) : + llama_model_deepseek4::graph(params) { + const auto & model = static_cast(model_base); + + topk_of.assign(n_layer, nullptr); + + ggml_tensor * cur; + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + llm_graph_input_dsv4 * inp_dsv4 = build_inp_dsv4(); + ggml_build_forward_expand(gf, inp_dsv4->get_raw()->self_kq_mask); + + // host-side inputs + auto inp41 = std::make_unique(model, inp_dsv4->mctx); + std::vector engram_ids(model.engram_layers.size(), nullptr); + for (size_t e = 0; e < engram_ids.size(); ++e) { + engram_ids[e] = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t) model.engram_n_cols()*n_tokens); + ggml_set_input(engram_ids[e]); + ggml_set_name(engram_ids[e], ("engram_ids-" + std::to_string(model.engram_layers[e])).c_str()); + } + inp41->engram_ids = engram_ids; + ggml_tensor * engram_keep = nullptr; + if (!engram_ids.empty() && ubatch.embd) { + engram_keep = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens); + ggml_set_input(engram_keep); + ggml_set_name(engram_keep, "engram_keep"); + inp41->engram_keep = engram_keep; + } + ggml_tensor * cand_pin = nullptr; + if (hparams.dsv41_cand_src >= 0) { + const int32_t cs = hparams.dsv41_cand_src; + inp41->cand_csa = hparams.dsv4_compress_ratios[cs] == 2; + const auto & ci = inp41->cand_csa ? inp_dsv4->get_csa() : inp_dsv4->get_hca(); + if (ci.kq_mask) { + const int64_t bs = hparams.dsv41_cand_block_size; + cand_pin = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, (ci.kq_mask->ne[0] + bs - 1)/bs, + ci.kq_mask->ne[1], 1, ci.kq_mask->ne[3]); + ggml_set_input(cand_pin); + ggml_set_name(cand_pin, "cand_pin"); + inp41->cand_pin = cand_pin; + } + } + res->add_input(std::move(inp41)); + + const int64_t hc = hparams.dsv4_hc_mult; + ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens); + inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); + cb(inpL, "hc_init", -1); + + // single-pass mHC: the pre-mix used to collapse the stream lags one sublayer + ggml_tensor * pre_prev = nullptr; + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + const int e = model.engram_index(il); + if (e >= 0) { + inpL = build_engram(model, inpL, engram_ids[e], engram_keep, il); + cb(inpL, "engram_out", il); + } + + if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) { + res->t_layer_inp[il] = build_hc_mean(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; + + ggml_tensor * attn_pre = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, il); + + if (pre_prev == nullptr) { + // identity pre-mix of the first layer: copy 0 + cur = ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0)); + } else { + cur = build_hc_pre(inpL, pre_prev, 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_attention41(model, inp_dsv4, cur, inp_pos, cand_pin, il); + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + residual = inpL; + ggml_tensor * ffn_pre = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, il); + + cur = build_hc_pre(inpL, attn_pre, il); + cb(cur, "hc_ffn_pre", il); + + ggml_build_forward_expand(gf, residual); + ggml_build_forward_expand(gf, post); + ggml_build_forward_expand(gf, comb); + + cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + ggml_tensor * exp_probs_b = layer.ffn_exp_probs_b; + ASSERT_EMBD_OR_TOKEN(ubatch); + if (ubatch.embd != nullptr && layer.ffn_exp_probs_b_vl) { + exp_probs_b = layer.ffn_exp_probs_b_vl; + } + + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + 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); + inpL = build_cvec(inpL, il); + cb(inpL, "l_last", il); + + // flat [n_embd*hc, n_tokens] alias for parity tooling (compare_ref.py reads l_out-) + ggml_tensor * l_out = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + cb(l_out, "l_out", il); + ggml_build_forward_expand(gf, l_out); + + pre_prev = ffn_pre; + } + + if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) { + res->t_layer_inp[n_layer] = build_hc_mean(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); + if (inp_out_ids) { + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + pre_prev = ggml_get_rows(ctx0, pre_prev, inp_out_ids); + } + res->t_h_nextn = flat; + + // no hc_head: collapse with the last FFN's pre-mix + inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, flat->ne[1]); + cur = build_hc_pre(inpL, pre_prev, -1); + cb(cur, "hc_out", -1); + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + // llama_model_n_embd_out() is hc_mult*n_embd (the MTP / DSpark heads read the flat HC state through it), and + // embedding extraction / pooling copy n_embd_out floats per row: an n_embd-wide result_norm here read past the + // tensor end (GGML_ASSERT in ggml_backend_tensor_get). Embedding contexts therefore return the flat HC state + // [hc_mult*n_embd] per token (pooled as usual); logits are unaffected. + res->t_embd = cparams.embeddings ? flat : cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; +} diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 6f28d2f1db4c..7045cbd3c44f 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -81,6 +81,7 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { 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); + ml.get_key(LLM_KV_HYPER_CONNECTION_SINGLE_PASS, hparams.dsv4_hc_single_pass, false); GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero @@ -94,6 +95,12 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { } GGML_ASSERT(hparams.n_swa > 0); + if (hparams.dsv4_hc_single_pass) { + // DeepSeek-V4.1 DSpark: every block query sees the whole main-KV ring (the window_size positions up to the + // anchor's predecessor) plus the block. The graph gives all queries of a block the mask row of its first + // query, whose standard-SWA window must then reach back window_size positions before the anchor. + hparams.n_swa += 1; + } hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; hparams.set_swa_pattern(0); for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { @@ -256,9 +263,12 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t hc_dim = hc_mult * n_embd; const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult; - 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); + // single-pass mHC (DeepSeek-V4.1) collapses with the last FFN pre-mix: no hc_head + if (!hparams.dsv4_hc_single_pass) { + 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]; @@ -1279,12 +1289,265 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_ } } +// DeepSeek-V4.1 DSpark decoder (mtp.0..2 of the V4.1 checkpoint, model.py DSparkBlock / forward_spec), dual-mode +// by batch type like graph_dsv4: +// * embd batch -> main_x = main_norm(main_proj(mean-hc attention inputs of the target layers)), every stage's +// kv_norm(wkv(main_x)) roped at the target positions is written into its window ring +// * token batch -> [anchor, noise...] block through the stages; V4.1 blocks use single-pass mHC (attention +// collapses with the previous sublayer's pre-mix, identity copy 0 for the first stage, the FFN with +// this attention's), no per-head q norm, plain rope; collapse with the last FFN pre-mix (no +// hc_head); markov + confidence heads as V4 +llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const llm_graph_params & params) : + llama_model_deepseek41::graph(params) { + const int64_t n_embd_inp = hparams.n_embd_inp_enc(); + 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_ASSERT(hparams.dsv4_hc_single_pass); + + ggml_tensor * inp_pos = build_inp_pos(); + + llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa(); + + // stages are ratio-0 layers: plain rope at rope_theta, no YaRN (model.py precompute_freqs_cis, original_seq_len 0) + const auto rope = [&](ggml_tensor * t) { + t = ggml_rope_ext(ctx0, t, inp_pos, nullptr, n_embd_head_rope, rope_type, 0, + freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + return ggml_rope_set_offset(t, n_embd_head_nope); + }; + + ASSERT_EMBD_OR_TOKEN(ubatch); + if (ubatch.embd) { + auto inp = std::make_unique(n_embd_inp); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * inp_target = inp->embd; + cb(inp_target, "inp_target_features", -1); + + res->add_input(std::move(inp)); + + ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s); + inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1); + cb(inp_g, "main_x", -1); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + 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); + kv = rope(kv); + 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; + } + + 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(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; + const int64_t n_groups = hparams.dsv4_o_group_count; + const int64_t n_heads_group = n_head / n_groups; + const int64_t o_lora_rank = hparams.dsv4_o_lora_rank; + const int64_t o_group_dim = n_heads_group*n_embd_head; + const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + + GGML_ASSERT(n_embd_head == n_embd_head_v); + GGML_ASSERT(n_head % n_groups == 0); + + 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); + + // window of a block = window of its first query (get_dspark_topk_idxs: ring + block for every query), see n_swa + { + ggml_tensor * m = inp_attn->self_kq_mask_swa_cnv; + const int64_t n_blocks = ubatch.n_seqs_unq; + const int64_t bd = n_blocks > 0 ? n_tokens / n_blocks : 0; + if (m && bd > 1 && n_tokens % n_blocks == 0 && m->ne[2] == 1 && m->ne[1] % bd == 0 && ggml_is_contiguous(m)) { + ggml_tensor * m4 = ggml_reshape_4d(ctx0, m, m->ne[0], bd, m->ne[1]/bd, m->ne[3]); + ggml_tensor * r0 = ggml_view_4d(ctx0, m4, m4->ne[0], 1, m4->ne[2], m4->ne[3], m4->nb[1], m4->nb[2], m4->nb[3], 0); + m4 = ggml_repeat(ctx0, r0, m4); + inp_attn->self_kq_mask_swa_cnv = ggml_reshape_4d(ctx0, m4, m->ne[0], m->ne[1], 1, m->ne[3]); + cb(inp_attn->self_kq_mask_swa_cnv, "dspark_block_mask", -1); + } + } + + ggml_tensor * pre_prev = nullptr; + + 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 * attn_pre = build_hc_mixes(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, il); + + ggml_tensor * cur; + if (pre_prev == nullptr) { + // identity pre-mix of the first stage (make_identity_pre_mix): copy 0 + cur = ggml_cont(ctx0, ggml_view_2d(ctx0, inpL, n_embd, n_tokens, inpL->nb[2], 0)); + } else { + cur = build_hc_pre(inpL, pre_prev, il); + } + cb(cur, "hc_attn_pre", il); + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // DSparkAttention: block queries see the main-KV ring plus every block key (non-causal draft context) + { + ggml_tensor * qr = build_lora_mm(layer.wq_a, cur); + qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(qr, "qr_norm", il); + + ggml_tensor * q = build_lora_mm(layer.wq_b, qr); + q = ggml_reshape_3d(ctx0, q, n_embd_head, n_head, n_tokens); + q = rope(q); + cb(q, "q", il); + + ggml_tensor * kv = build_lora_mm(layer.wkv, cur); + 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); + kv = rope(kv); + cb(kv, "kv", il); + + ggml_tensor * out = build_attn(inp_attn, + nullptr, nullptr, nullptr, + q, kv, kv, + nullptr, layer.attn_sinks, nullptr, + kq_scale, il); + cb(out, "attn_raw", il); + + out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, n_tokens); + out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, 0, + freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + out = ggml_rope_set_offset(out, n_embd_head_nope); + cb(out, "attn_derope", il); + + out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, n_tokens); + out = ggml_permute(ctx0, out, 0, 2, 1, 3); + ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out); + cb(oa, "attn_wo_a", il); + oa = ggml_permute(ctx0, oa, 0, 2, 1, 3); + oa = ggml_cont_2d(ctx0, oa, o_lora_rank*n_groups, n_tokens); + + cur = build_lora_mm(layer.wo_b, oa); + cb(cur, "attn_out", il); + } + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + residual = inpL; + ggml_tensor * ffn_pre = build_hc_mixes(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, il); + + cur = build_hc_pre(inpL, attn_pre, 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); + + ggml_tensor * l_out = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + cb(l_out, "l_out", il); + ggml_build_forward_expand(gf, l_out); + + pre_prev = ffn_pre; + } + + // no hc_head: collapse with the last FFN's pre-mix (forward_head: hc_pre(x, pre_mix)) + ggml_tensor * cur = build_hc_pre(inpL, pre_prev, -1); + cb(cur, "hc_out", -1); + + // confidence head input: the un-normed collapsed hidden state + res->t_embd = cur; + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + + auto * output = model.output; + auto * output_s = model.output_s; + 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; + output_s = model_other->output_s; + } + + cur = build_lora_mm(output, cur, output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + if (model.dspark_markov_w1) { + build_dspark_markov_head(*this, model, inp_tokens); + } +} + std::unique_ptr llama_model_dflash::build_arch_graph(const llm_graph_params & params) const { switch (params.gtype) { case LLM_GRAPH_TYPE_ENCODER: return std::make_unique>(*this, params); case LLM_GRAPH_TYPE_DEFAULT: case LLM_GRAPH_TYPE_DECODER: + if (hparams.dsv4_hc_mult > 0 && hparams.dsv4_hc_single_pass) { + return std::make_unique(*this, params); + } if (hparams.dsv4_hc_mult > 0) { return std::make_unique(*this, params); } diff --git a/src/models/models.h b/src/models/models.h index e6cf9d213cee..af6be6e8050a 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1424,6 +1424,76 @@ struct llama_model_deepseek4 : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +// DeepSeek-V4.1: V4 graph pieces (HC, MLA window attention, sparse compressed attention, MoE) plus Engram, +// single-pass mHC, ratio-1/2 compressors with cross-layer KV / index-K / top-k sharing and two-level candidate selection. +struct llama_model_deepseek41 : public llama_model_base { + llama_model_deepseek41(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + // Engram n-gram hash constants (host side; ggml has no int64 mul/xor/mod) + std::vector engram_layers; // backbone layer ids + std::vector engram_n_rows; // table rows per engram layer + uint32_t engram_max_ngram = 0; + uint32_t engram_n_head = 0; + uint32_t engram_head_dim = 0; + int64_t engram_pad_cid = 0; // compressed id of the pad token + std::vector engram_token_map; // raw id -> compressed id + std::vector engram_mult; // [n_engram][max_ngram] + std::vector engram_primes; // [n_engram][(max_ngram-1)*n_head] + std::vector engram_offsets; // [n_engram][(max_ngram-1)*n_head] + + int engram_index(int il) const { + for (size_t i = 0; i < engram_layers.size(); ++i) { + if (engram_layers[i] == il) { + return (int) i; + } + } + return -1; + } + uint32_t engram_n_cols() const { return (engram_max_ngram - 1)*engram_n_head; } + + // hash ids of every token of a ubatch for engram layer e: out[i*n_cols + c] + void engram_hash(int e, const llama_ubatch & ubatch, const std::vector & prev, std::vector & out) const; + + struct graph : public llama_model_deepseek4::graph { + graph(const llama_model & model, const llm_graph_params & params); + // helpers only (no graph built): base of the DSpark draft graph (llama_model_dflash::graph_dsv41) + explicit graph(const llm_graph_params & params) : llama_model_deepseek4::graph(params) {} + + // HC coefficients of one sublayer: returns pre [hc, n_tokens], fills post / comb + ggml_tensor * build_hc_mixes( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) const; + + ggml_tensor * build_engram( + const llama_model_deepseek41 & model, + ggml_tensor * x, + ggml_tensor * ids, + ggml_tensor * keep, + int il) const; + + ggml_tensor * build_attention41( + const llama_model_deepseek41 & model, + llm_graph_input_dsv4 * inp_dsv4, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * cand_pin, + int il); + + // per-graph sharing state (layers run in order, sources before consumers) + std::vector topk_of; // [n_layer] top-k of each index source + ggml_tensor * cand_mask = nullptr; // candidate-block mask from the candidate source layer + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_deepseek2ocr : public llama_model_base { llama_model_deepseek2ocr(const struct llama_model_params & params) : llama_model_base(params) {} @@ -1484,6 +1554,11 @@ struct llama_model_dflash : public llama_model_base { graph_dsv4(const llama_model & model, const llm_graph_params & params); }; + // DeepSeek-V4.1 DSpark: single-pass mHC (lagged pre-mix, no hc_head), no per-head q norm + struct graph_dsv41 : public llama_model_deepseek41::graph { + graph_dsv41(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 20b979ed0758..785d1600ed40 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5192,6 +5192,37 @@ static bool graph_mul_mat_hi_prec_act(ggml_cgraph * gf, ggml_op op) { return false; } +// GGML_OP_MUL_MAT with src1 requested at GGML_PREC_F32 ([TAG_GGML_PREC]): F32 weights must not be rounded to a +// lower precision (e.g. TF32 tensor cores) at any batch size, so the tolerance is at the F32 level, not 5e-4 +struct test_mul_mat_prec_f32 : public test_case { + const int64_t m; + const int64_t n; + const int64_t k; + const int64_t bs; + + std::string vars() override { + return VARS_TO_STR4(m, n, k, bs); + } + + double max_nmse_err() override { + return 1e-10; // F32 accumulation-order noise is ~1e-14; TF32 inputs give ~1e-7 + } + + test_mul_mat_prec_f32(int64_t m = 128, int64_t n = 7, int64_t k = 1024, int64_t bs = 1) + : m(m), n(n), k(k), bs(bs) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, m, bs); + ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, n, bs); + ggml_set_name(a, "a"); + ggml_set_name(b, "b"); + ggml_tensor * out = ggml_mul_mat(ctx, a, b); + ggml_prec_set_src(out, GGML_PREC_F32, 1); + ggml_set_name(out, "out"); + return out; + } +}; + // GGML_OP_MUL_MAT struct test_mul_mat : public test_case { const ggml_type type_a; @@ -9829,6 +9860,16 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_get_rows(GGML_TYPE_F32, 256, 8, 2, 1, 1, false, true, 3)); + // MXFP8 (E4M3 + E8M0, 33 B / 32): gather-only type for large host-resident tables + for (int b : {1, 7}) { + for (bool v : {false, true}) { + for (bool vs0 : {false, true}) { + test_cases.emplace_back(new test_get_rows(GGML_TYPE_MXFP8, 256, 5, 4, b, 1, v, vs0)); + } + } + } + test_cases.emplace_back(new test_get_rows(GGML_TYPE_MXFP8, 256, 1000, 24, 1, 1, false)); + test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 8, 2, 1, false)); test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 70000, 4, 1, false)); // row count > CUDA grid-y limit (65535) for (ggml_type type : all_types) { @@ -10757,6 +10798,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); + // src1 at GGML_PREC_F32: no TF32 at any batch size (mmvf / mmf / cuBLAS ranges), incl. a batched (attention) shape + for (int64_t n : {1, 3, 4, 7, 8, 9, 16, 17, 64}) { + test_cases.emplace_back(new test_mul_mat_prec_f32(384, n, 5120, 1)); + test_cases.emplace_back(new test_mul_mat_prec_f32(256, n, 512, 4)); + } + // m == 1, with n on both sides of MMVF_MAX_BATCH_SIZE (8): mmvf below, operand swap above for (int64_t n : {1, 7, 8, 9, 16, 127, 128, 511, 512}) { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 1, n, 2048, {1, 1}, {1, 1})); diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 49519eecab1d..bd38e488dbfd 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -4381,6 +4381,146 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + // DeepSeek V4.1 tests - V4 behaviour, DSML tag names with a leading space: + // <｜DSML｜ calls> / <｜DSML｜ invoke name="..."> / <｜DSML｜ parameter name="..." string="..."> + { + auto tst = peg_tester("models/templates/deepseek-ai-DeepSeek-V4.1-Flash-1017.jinja", detailed_debug); + + // Pure content (non-thinking mode; generation prompt ends with ) + tst.test("Hello, world!\nWhat's up?") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect(message_assist) + .run(); + + // Thinking + content + tst.test("I'm\nthinkingHello, world!\nWhat's up?") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect(message_assist_thoughts) + .run(); + + // Thinking + tool call (single, string param) + tst.test( + "Let me check the time\n\n" + "<｜DSML｜ calls>\n" + "<｜DSML｜ invoke name=\"get_time\">\n" + "<｜DSML｜ parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time")) + .run(); + + // Tool call without reasoning (non-thinking mode), integer param (string="false") + tst.test( + "<｜DSML｜ calls>\n" + "<｜DSML｜ invoke name=\"special_function\">\n" + "<｜DSML｜ parameter name=\"arg1\" string=\"false\">1\n" + "\n" + "") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ special_function_tool }) + .expect(message_assist_call) + .run(); + + // Multiple parallel tool calls with reasoning + tst.test( + "Calling both\n\n" + "<｜DSML｜ calls>\n" + "<｜DSML｜ invoke name=\"get_time\">\n" + "<｜DSML｜ parameter name=\"city\" string=\"true\">Paris\n" + "\n" + "<｜DSML｜ invoke name=\"get_weather\">\n" + "<｜DSML｜ parameter name=\"city\" string=\"true\">Paris\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .parallel_tool_calls(true) + .tools({ get_time_tool, get_weather_tool }) + .expect(message_with_reasoning_content_and_multiple_tool_calls( + "Calling both", "", + { { "get_time", R"({"city": "Paris"})" }, { "get_weather", R"({"city": "Paris"})" } })) + .run(); + + // Tool call with content before tool calls + tst.test( + "Thinking about it" + "Let me call the function.\n\n" + "<｜DSML｜ calls>\n" + "<｜DSML｜ invoke name=\"special_function\">\n" + "<｜DSML｜ parameter name=\"arg1\" string=\"false\">1\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ special_function_tool }) + .expect_reasoning("Thinking about it") + .expect_content("Let me call the function.") + .expect_tool_calls({ + { "special_function", R"({"arg1": 1})", {} }, + }) + .expect_reconstruction() + .run(); + + // Tool call with multiple params (mixed types) + tst.test( + "Multi-arg call\n\n" + "<｜DSML｜ calls>\n" + "<｜DSML｜ invoke name=\"magic_int\">\n" + "<｜DSML｜ parameter name=\"ref\" string=\"false\">42\n" + "<｜DSML｜ parameter name=\"name\" string=\"true\">foo bar\n" + "\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ magic_int_tool }) + .expect_reasoning("Multi-arg call") + .expect_tool_calls({ + { "magic_int", R"({"ref": 42, "name": "foo bar"})", {} }, + }) + .run(); + + // Continuation tests + tst.test("world!\nWhat's up?") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .enable_thinking(true) + .messages({ message_user, message_assist_prefill_content }) + .add_generation_prompt(false) + .continue_final_message(COMMON_CHAT_CONTINUATION_CONTENT) + .expect_reasoning("I'm thinking") + .expect_content("Hello, world!\nWhat's up?") + .run(); + + tst.test(" thinkingHello, world!\nWhat's up?") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .enable_thinking(true) + .messages({ message_user, message_assist_prefill_reasoning }) + .add_generation_prompt(false) + .continue_final_message(COMMON_CHAT_CONTINUATION_REASONING) + .expect_reasoning("I'm thinking") + .expect_content("Hello, world!\nWhat's up?") + .run(); + + tst.test( + "Let me check the time\n\n" + "<｜DSML｜ calls>\n" + "<｜DSML｜ invoke name=\"get_time\">\n" + "<｜DSML｜ parameter name=\"city\" string=\"true\">Tokyo\n" + "\n" + "") // no after the TC close because the grammar will immediately constrain it to end + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ get_time_tool }) + .expect_reasoning("Let me check the time") + .expect_tool_calls({ { "get_time", R"({"city": "Tokyo"})", {} } }) + .run(); + } + { // The DSML separator belongs to the tool call block, not assistant content. auto tst = peg_tester("models/templates/deepseek-ai-DeepSeek-V4-Flash-0731.jinja", detailed_debug); @@ -7825,6 +7965,74 @@ static void test_deepseek_v4_tool_result_ordering() { } } +// DeepSeek V4.1: non-OpenAI message keys reach the template (encoding.py golden cases 3 and 4) and +// content blocks around media markers are joined with "\n\n" like encoding.py does around images. +static void test_deepseek_v4_1_message_fields() { + LOG_DBG("%s\n", __func__); + + auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V4.1-Flash-1017.jinja"); + auto render = [&](const char * messages_json, bool add_generation_prompt) { + common_chat_templates_inputs inputs; + inputs.messages = common_chat_msgs_parse_oaicompat(json::parse(messages_json)); + inputs.add_generation_prompt = add_generation_prompt; + inputs.chat_template_kwargs["thinking_mode"] = "\"chat\""; + return common_chat_templates_apply(tmpls.get(), inputs).prompt.str(); + }; + + // "task" on a user turn selects the task token + assert_equals(std::string("<｜begin▁of▁sentence｜><｜User｜>q<｜Assistant｜><｜action｜>"), + render(R"([{"role": "user", "content": "q", "task": "action"}])", false)); + + // per-message "tools" on a later system message, tool call "namespace" + { + auto prompt = render(R"([ + {"role": "system", "content": "a"}, + {"role": "system", "content": "b", "tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}]}, + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "", "tool_calls": [{"type": "function", "namespace": "ns", "function": {"name": "f", "arguments": {"x": 1}}}]} + ])", false); + assert_contains(prompt, "<｜System｜>b\n\n## Tools"); + assert_contains(prompt, "{\"name\": \"f\", \"parameters\": {\"type\": \"object\"}}"); + assert_contains(prompt, "<｜DSML｜ invoke name=\"ns::f\">"); + } + + // media markers stay separate blocks + assert_equals(std::string("<｜begin▁of▁sentence｜><｜User｜>a\n\n<__media__>\n\nb<｜Assistant｜>"), + render(R"([{"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "media_marker", "text": "<__media__>"}, {"type": "text", "text": "b"}]}])", true)); + + // OpenAI reasoning_effort values never fail: minimal -> low, medium -> unset (template default high), xhigh -> max + for (const auto & [effort, budget] : std::vector>{ + { "minimal", "50" }, { "low", "50" }, { "medium", "75" }, { "high", "75" }, { "xhigh", "100" }, { "max", "100" } }) { + common_chat_templates_inputs inputs; + inputs.messages = { message_user }; + inputs.add_generation_prompt = true; + inputs.chat_template_kwargs["reasoning_effort"] = json(effort).dump(); + auto prompt = common_chat_templates_apply(tmpls.get(), inputs).prompt.str(); + assert_contains(prompt, "<｜System｜>Reasoning Effort: " + budget + " (range 1-100"); + } + + // tools declared only on a system message are still parsed as tool calls + { + common_chat_templates_inputs inputs; + inputs.messages = common_chat_msgs_parse_oaicompat(json::parse(R"([ + {"role": "system", "content": "s", "tools": [{"type": "function", "function": {"name": "get_time", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]}, + {"role": "user", "content": "q"} + ])")); + inputs.add_generation_prompt = true; + inputs.chat_template_kwargs["thinking_mode"] = "\"chat\""; + inputs.enable_thinking = false; + inputs.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK; + common_chat_session session(tmpls.get(), nullptr, inputs, {}); + const auto & msg = session.finish(common_chat_input(std::string( + "\n\n<｜DSML｜ calls>\n<｜DSML｜ invoke name=\"get_time\">\n" + "<｜DSML｜ parameter name=\"city\" string=\"true\">Tokyo\n\n"))); + assert_equals(size_t(1), msg.tool_calls.size()); + assert_equals(std::string("get_time"), msg.tool_calls[0].name); + assert_equals(std::string("{\"city\":\"Tokyo\"}"), msg.tool_calls[0].arguments); + } +} + static void test_prompt_parts() { LOG_DBG("%s\n", __func__); const std::string tmpl_src = @@ -8175,6 +8383,7 @@ int main(int argc, char ** argv) { test_translate_gemma(); test_deepseek_v4_thinking_retention(); test_deepseek_v4_tool_result_ordering(); + test_deepseek_v4_1_message_fields(); test_template_generation_prompt(); test_reasoning_effort_caps(); test_prompt_parts(); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 94a22ff2f57f..693db013f85e 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -187,6 +187,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { n_head = 8; n_ff = 1024; n_layer = 4; + } else if (arch == LLM_ARCH_DEEPSEEK41) { + // window-only, Engram, ratio-2 and ratio-1 sources, top-k reuse and a candidate-masked index layer + n_embd = 512; + n_head = 8; + n_ff = 1024; + n_layer = 6; } else if (arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_LAGUNA) { n_embd = 160; // exercise per-head tensor split granularity with head size 80 } else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) { @@ -277,11 +283,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer); } else { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); - ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head_kv); + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) ? uint32_t(1) : n_head_kv); } ms.add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS, 8.0f); - if (arch == LLM_ARCH_DEEPSEEK4) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) { ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, n_embd_head); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, n_embd_head); ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, n_embd_head/2); @@ -326,7 +332,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, 1e-5f); ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS, 1e-5f); ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS, uint32_t(8)); - ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(64) : uint32_t(512)); + ms.add_kv(LLM_KV_ATTENTION_Q_LORA_RANK, (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) ? uint32_t(64) : uint32_t(512)); ms.add_kv(LLM_KV_ATTENTION_KV_LORA_RANK, uint32_t(512)); ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8)); ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW, n_ctx/8); @@ -462,6 +468,41 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } + if (arch == LLM_ARCH_DEEPSEEK41) { + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8)); + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32)); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector({0, 0, 2, 2, 1, 1})); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f); + ms.add_kv(LLM_KV_ATTENTION_KV_SOURCE_LAYERS, std::vector({2, 4})); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_SOURCE_LAYERS, std::vector({2, 4, 5})); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_SOURCE_LAYER, int32_t(4)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_TOP_K_BLOCKS, uint32_t(4096)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_CANDIDATE_BLOCK_SIZE, uint32_t(4)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); + ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); + // Engram on layers 1 and 3: 2..3-grams, 2 heads, 64-wide rows + std::vector token_map(n_vocab); + for (uint32_t i = 0; i < n_vocab; ++i) { + token_map[i] = (int32_t) (i % 61); + } + ms.add_kv(LLM_KV_ENGRAM_LAYERS, std::vector({1, 3})); + ms.add_kv(LLM_KV_ENGRAM_NUM_EMBEDDINGS, std::vector({420, 508})); + ms.add_kv(LLM_KV_ENGRAM_MAX_NGRAM_SIZE, uint32_t(3)); + ms.add_kv(LLM_KV_ENGRAM_VOCAB_SIZE, uint32_t(100)); + ms.add_kv(LLM_KV_ENGRAM_HEAD_COUNT, uint32_t(2)); + ms.add_kv(LLM_KV_ENGRAM_HEAD_DIM, uint32_t(64)); + ms.add_kv(LLM_KV_ENGRAM_PAD_TOKEN_ID, uint32_t(2)); + ms.add_kv(LLM_KV_ENGRAM_COMPRESSED_VOCAB_SIZE, uint32_t(61)); + ms.add_kv(LLM_KV_ENGRAM_TOKEN_MAP, token_map); + ms.add_kv(LLM_KV_ENGRAM_MULTIPLIERS, std::vector({ 76632096046245ull, 4839876093313ull, 35959672319349ull, + 67716810739261ull, 51510806800915ull, 30921347202721ull })); + ms.add_kv(LLM_KV_ENGRAM_PRIMES, std::vector({101, 103, 107, 109, 113, 127, 131, 137})); + } + if (arch == LLM_ARCH_HRM_TEXT) { // 6 cache slots alias 2 physical blocks: 1 low-stack layer + 1 high-stack layer ms.add_kv(LLM_KV_HRM_LAYERS_PER_STACK, uint32_t(1)); @@ -505,6 +546,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { gating_func = 1; break; case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: gating_func = 4; break; default: @@ -762,6 +804,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DOTS3NOTE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_DEEPSEEK41: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: case LLM_ARCH_EXAONE_MOE: