diff --git a/common/arg.cpp b/common/arg.cpp index 183293c8bed..d844c436094 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -377,6 +377,10 @@ common_models_handler common_models_handler_init(const common_params & params, l params.speculative.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3) != params.speculative.types.end(); + const bool spec_type_draft_dspark = std::find(params.speculative.types.begin(), + params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) != params.speculative.types.end(); + // only download mmproj if the current example is using it bool use_mmproj = false; for (const auto & ex : mmproj_examples) { @@ -391,6 +395,7 @@ common_models_handler common_models_handler_init(const common_params & params, l opts.download_mtp = spec_type_draft_mtp; opts.download_eagle3 = spec_type_draft_eagle3; opts.download_dflash = spec_type_draft_dflash; + opts.download_dspark = spec_type_draft_dspark; opts.download_mmproj = use_mmproj && !params.no_mmproj && params.mmproj.path.empty() && params.mmproj.url.empty(); @@ -405,6 +410,7 @@ common_models_handler common_models_handler_init(const common_params & params, l opts_spec.download_mtp = true; opts_spec.download_dflash = true; opts_spec.download_eagle3 = true; + opts_spec.download_dspark = true; } plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts_spec); } @@ -547,12 +553,19 @@ void common_models_handler_apply(common_models_handler & handler, common_params plan_spec.mtp = {}; plan_spec.dflash = {}; plan_spec.eagle3 = {}; + plan_spec.dspark = {}; } // infer the speculative type from the sidecar shipped by the draft repo when none is requested if (spec_types_is_default(params)) { if (!plan_spec.mtp.local_path.empty()) { params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP }; + plan_spec.dspark = {}; + plan_spec.dflash = {}; + plan_spec.eagle3 = {}; + } else if (!plan_spec.dspark.local_path.empty()) { + // dspark outranks dflash, its sidecar carries the extra Markov head + params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK }; plan_spec.dflash = {}; plan_spec.eagle3 = {}; } else if (!plan_spec.dflash.local_path.empty()) { @@ -566,7 +579,8 @@ void common_models_handler_apply(common_models_handler & handler, common_params // when a sidecar type is requested, the draft repo resolves to its sidecar instead of a full model const bool spec_sidecar_found = !plan_spec.mtp.local_path.empty() || !plan_spec.dflash.local_path.empty() || - !plan_spec.eagle3.local_path.empty(); + !plan_spec.eagle3.local_path.empty() || + !plan_spec.dspark.local_path.empty(); if (!plan_spec.mtp.local_path.empty() && !had_spec_url) { tasks.emplace_back(plan_spec.mtp, opts, [&]() { // only use the discovered MTP head when no draft path is set yet @@ -597,6 +611,16 @@ void common_models_handler_apply(common_models_handler & handler, common_params } }); } + if (!plan_spec.dspark.local_path.empty() && !had_spec_url) { + tasks.emplace_back(plan_spec.dspark, opts, [&]() { + // only use the discovered DSpark sidecar when no draft path is set yet + if (params.speculative.draft.mparams.path.empty()) { + params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dspark); + } else { + hf_cache::finalize_file(plan_spec.dspark); + } + }); + } // a wired draft sidecar counts as an explicit draft for the main plan fallback below if (spec_sidecar_found) { @@ -652,6 +676,16 @@ void common_models_handler_apply(common_models_handler & handler, common_params } }); } + if (!plan.dspark.local_path.empty() && !had_spec_url) { + tasks.emplace_back(plan.dspark, opts, [&]() { + // only fall back to the discovered DSpark sidecar when no draft was explicitly provided + if (params.speculative.draft.mparams.empty()) { + params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.dspark); + } else { + hf_cache::finalize_file(plan.dspark); + } + }); + } if (!plan.preset.local_path.empty()) { tasks.emplace_back(plan.preset, opts, [&]() { // if HF repo is a preset repo, we simply run server in router mode with the preset.ini file diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp index f786f5ff231..1910b4f1e13 100644 --- a/common/chat-peg-parser.cpp +++ b/common/chat-peg-parser.cpp @@ -6,6 +6,9 @@ #include +#include +#include + using ordered_json = nlohmann::ordered_json; static std::string_view trim_trailing_space(std::string_view sv, int max = -1) { @@ -235,6 +238,43 @@ common_peg_parser common_chat_peg_builder::tag_with_safe_content(const std::stri return zero_or_more(choice({ p, content_chunk })); } +common_peg_parser common_chat_peg_builder::permute(const std::string & rule_prefix, + const std::vector & parsers) { + if (parsers.empty()) { + return eps(); + } + + if (parsers.size() == 1 || parsers.size() > COMMON_CHAT_MAX_PERMUTE) { + return sequence(parsers); + } + + std::map rules; + std::function remaining_of; + + remaining_of = [&](uint32_t remaining) -> common_peg_parser { + if (remaining == 0) { + return eps(); + } + + auto cached = rules.find(remaining); + if (cached != rules.end()) { + return cached->second; + } + + auto alternatives = choice(); + for (size_t i = 0; i < parsers.size(); i++) { + const uint32_t bit = 1u << i; + if (remaining & bit) { + alternatives |= parsers[i] + remaining_of(remaining & ~bit); + } + } + + return rules.emplace(remaining, rule(rule_prefix + "-" + std::to_string(remaining), alternatives)).first->second; + }; + + return remaining_of((1u << parsers.size()) - 1); +} + std::string & common_chat_peg_mapper::args_target() { return (current_tool && !current_tool->name.empty()) ? current_tool->arguments : args_buffer; } diff --git a/common/chat-peg-parser.h b/common/chat-peg-parser.h index cd14f2c1175..5d764dbaa0e 100644 --- a/common/chat-peg-parser.h +++ b/common/chat-peg-parser.h @@ -55,6 +55,8 @@ class common_chat_peg_minimax_m3_mapper : public common_chat_peg_mapper { struct content_structure; struct tool_call_structure; +constexpr size_t COMMON_CHAT_MAX_PERMUTE = 6; + class common_chat_peg_builder : public common_peg_parser_builder { public: // Tag constants (from former common_chat_peg_base_builder) @@ -105,6 +107,9 @@ class common_chat_peg_builder : public common_peg_parser_builder { common_peg_parser tool_arg_json_value(const common_peg_parser & p) { return tag(TOOL_ARG_VALUE, p); } + // Matches every parser exactly once, in any order. + common_peg_parser permute(const std::string & rule_prefix, const std::vector & parsers); + // Return a parser that parses the prefix of a string, up to a given delimiter. common_peg_parser prefix(const std::string & s, const std::string & delimiter = {}); diff --git a/common/chat.cpp b/common/chat.cpp index 087a7fcfab4..e252bd50824 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1110,6 +1110,172 @@ static common_chat_params common_chat_params_init_ministral_3(const common_chat_ return data; } +static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + const std::string GEN_PREFIX = "<|im_start|>assistant\n"; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + + auto supports_reasoning = tmpl.source().find("") != std::string::npos; + + data.supports_thinking = supports_reasoning; + data.preserved_tokens = { + "", + "", + }; + + if (supports_reasoning) { + data.thinking_start_tag = ""; + // Support both and as reasoning end sequences. + // ", "" }; + data.preserved_tokens.insert(data.preserved_tokens.end(), { "", "" }); + } + + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" }, + { COMMON_CHAT_ROLE_TOOL, "<|im_start|>user\n" }, // Qwen3-Coder, Qwen3.5, Nemotron Nano 3 + { COMMON_CHAT_ROLE_TOOL, "<|im_start|>tool_response" }, // StepFun-3.5-Flash + { COMMON_CHAT_ROLE_USER, "<|im_start|>user" }, + { COMMON_CHAT_ROLE_SYSTEM, "<|im_start|>system" }, + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty(); + 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); + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = GEN_PREFIX; + if (supports_reasoning) { + data.generation_prompt += "\n" + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += "\n\n\n"; + } + } + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += msg.render_content(); + } + + data.prompt += data.generation_prompt; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.literal(GEN_PREFIX); + + auto reasoning = p.eps(); + if (supports_reasoning && extract_reasoning) { + reasoning = p.optional("" + p.space() + + p.reasoning(p.until_one_of({ "", "" })) + + (p.literal("") | p.peek(p.literal("")))); + } + + // Response format parser + if (has_response_format) { + return generation_prompt + (reasoning << p.content(p.schema(p.json(), "response-format", inputs.json_schema))); + } + + // Tool call parser + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + auto arg_close = p.tool_arg_close(p.literal("\n\n")); + auto arg_string = p.rule("xml-arg-string", + p.ac(p.tool_arg_string_value(p.until("\n\n")) + arg_close, "\n\n")); + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + auto parameters = function.contains("parameters") ? function.at("parameters") : json::object(); + + auto schema_info = common_schema_info(); + schema_info.resolve_refs(parameters); + + std::vector required_args; + std::vector optional_args; + + foreach_parameter(function, [&](const std::string & param_name, const json & param_schema, bool is_required) { + auto rule_name = "tool-" + name + "-arg-" + param_name; + + auto arg_open = p.tool_arg_open("\n"); + + auto arg_value = schema_info.resolves_to_string(param_schema) ? + arg_string : + p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", param_schema)) + arg_close; + + auto arg_rule = p.rule(rule_name, p.tool_arg(arg_open + arg_value)); + + (is_required ? required_args : optional_args).push_back(arg_rule); + }); + + // Accept required arguments in any order, as Qwen does not always adhere to the + // order provided. + auto args = p.permute("tool-" + name + "-args", required_args); + if (!optional_args.empty()) { + args = args + p.zero_or_more(p.choice(optional_args)); + } + + auto func = p.tool(p.tool_open("\n") + + p.tool_args(args) + + p.tool_close(p.literal("\n"))); + + tool_choice |= p.rule("tool-" + name, func); + }); + + auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0; + + // Qwen3-Coder models may occasionally omit the token. + auto tool_call_body = tool_choice + "" + p.space(); + auto tool_call_first = p.rule("tool-call-first", p.optional(p.literal("\n")) + tool_call_body); + auto tool_call = p.rule("tool-call", "\n" + tool_call_body); + + auto calls = inputs.parallel_tool_calls ? tool_call_first + p.zero_or_more(tool_call) : tool_call_first; + auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(calls, min_calls, 1)); + + return generation_prompt + + (reasoning << p.content(p.until_one_of({ "", "" }, + // Trigger on " common_chat_try_specialized_template( return common_chat_params_init_minicpm5(tmpl, params); } + // Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash + if (src.find("") != std::string::npos && + src.find(" common_list_cached_models() { split.prefix.find("mmproj") != std::string::npos || split.prefix.find("mtp-") != std::string::npos || split.prefix.find("eagle3-") != std::string::npos || - split.prefix.find("dflash-") != std::string::npos) { + split.prefix.find("dflash-") != std::string::npos || + split.prefix.find("dspark-") != std::string::npos) { continue; } if (seen.insert(f.repo_id + ":" + split.tag).second) { diff --git a/common/download.h b/common/download.h index 3e789e9e936..9a03f5e9147 100644 --- a/common/download.h +++ b/common/download.h @@ -59,6 +59,7 @@ struct common_download_opts { bool download_mtp = false; bool download_eagle3 = false; bool download_dflash = false; + bool download_dspark = false; common_download_callback * callback = nullptr; }; @@ -110,6 +111,7 @@ struct common_download_hf_plan { hf_cache::hf_file mtp; hf_cache::hf_file eagle3; hf_cache::hf_file dflash; + hf_cache::hf_file dspark; hf_cache::hf_file preset; // if set, only this file is downloaded }; common_download_hf_plan common_download_get_hf_plan(const common_params_model & model, const common_download_opts & opts); diff --git a/common/speculative.cpp b/common/speculative.cpp index 1f3263a3016..3e0e9f87295 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1336,7 +1336,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { GGML_ASSERT(ctx_tgt && ctx_dft && "MTP requires ctx_tgt and ctx_dft to be set"); n_embd = llama_model_n_embd_out(llama_get_model(ctx_dft)); - GGML_ASSERT(n_embd == llama_model_n_embd(llama_get_model(ctx_tgt)) && + GGML_ASSERT(n_embd == llama_model_n_embd_out(llama_get_model(ctx_tgt)) && "MTP input row width must match the target h_nextn width"); n_mtp_layers = std::max(1, (int) llama_model_n_layer_nextn(llama_get_model(ctx_dft))); diff --git a/conversion/__init__.py b/conversion/__init__.py index 682ec63fb72..11121e054b1 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -56,6 +56,9 @@ "Qwen3DSparkModel": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DFlashLagunaForCausalLM": "laguna", + + "DeepseekV4DSparkModel": "deepseek", + "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", "DistilBertModel": "bert", diff --git a/conversion/deepseek.py b/conversion/deepseek.py index ea6ae23d58e..5b69e23437f 100644 --- a/conversion/deepseek.py +++ b/conversion/deepseek.py @@ -475,7 +475,10 @@ def set_gguf_parameters(self): @ModelBase.register("DeepseekV4ForCausalLM") class DeepseekV4Model(TextModel): model_arch = gguf.MODEL_ARCH.DEEPSEEK4 + supports_mtp_export = True _skipped_mtp_tensors = 0 + _dsv4_main_layers: int | None = None + _dsv4_nextn_layers: int = 0 def __init__(self, *args, **kwargs): type(self)._skipped_mtp_tensors = 0 @@ -487,6 +490,8 @@ def __init__(self, *args, **kwargs): self.hparams.setdefault(key, value) self.block_count = self.hparams["num_hidden_layers"] + if self.mtp_only: + self.block_count += self.hparams.get("num_nextn_predict_layers", 0) self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) self._dsv4_fp8_dequantized: set[str] = set() @@ -504,13 +509,63 @@ def __init__(self, *args, **kwargs): with open(template_path, "r", encoding="utf-8") as f: self.gguf_writer.add_chat_template(f.read()) + def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]: + type(self)._dsv4_main_layers = self.hparams["num_hidden_layers"] + type(self)._dsv4_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0) + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) + @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: - name, _ = item + name, gen = item if name.startswith("mtp."): - cls._skipped_mtp_tensors += 1 - return None - return super().filter_tensors(item) + if not cls.mtp_only: + cls._skipped_mtp_tensors += 1 + return None + + assert cls._dsv4_main_layers is not None + parts = name.split(".", 2) + if len(parts) < 3 or not parts[1].isdecimal(): + raise ValueError(f"Unexpected DeepSeek-V4 MTP tensor {name!r}") + + mtp_idx = int(parts[1]) + if mtp_idx >= cls._dsv4_nextn_layers: + raise ValueError(f"Unexpected DeepSeek-V4 MTP layer {mtp_idx}") + + bid = cls._dsv4_main_layers + mtp_idx + suffix = parts[2] + root_hc_head = { + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + } + if suffix in root_hc_head: + name = suffix + elif suffix in ( + "e_proj.weight", "e_proj.scale", + "h_proj.weight", "h_proj.scale", + ): + name = f"layers.{bid}.nextn.{suffix}" + elif suffix == "enorm.weight": + name = f"layers.{bid}.nextn.enorm.weight" + elif suffix == "hnorm.weight": + name = f"layers.{bid}.nextn.hnorm.weight" + elif suffix == "norm.weight": + name = f"layers.{bid}.nextn.shared_head_norm.weight" + else: + name = f"layers.{bid}.{suffix}" + return name, gen + + if cls.mtp_only: + keep = name in ( + "embed.weight", + "norm.weight", + "head.weight", + "head.scale", + ) + if not keep: + return None + + return super().filter_tensors((name, gen)) @staticmethod def _float8_dtypes() -> tuple[torch.dtype, ...]: @@ -565,6 +620,10 @@ def set_gguf_parameters(self): self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hparams["hc_sinkhorn_iters"]) self.gguf_writer.add_hyper_connection_epsilon(hparams["hc_eps"]) self.gguf_writer.add_hash_layer_count(hparams["num_hash_layers"]) + if self.model_arch == gguf.MODEL_ARCH.DEEPSEEK4: + self.gguf_writer.add_embedding_length_out(hparams["hidden_size"] * hparams["hc_mult"]) + if self.mtp_only and (num_nextn_predict_layers := hparams.get("num_nextn_predict_layers", 0)) > 0: + self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers) def dequant_model(self): fp8_dtypes = self._float8_dtypes() @@ -669,12 +728,37 @@ def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: if self._dsv4_mxfp4_generated: return () - consumed: list[str] = self._write_hash_routing_tensors() + consumed: list[str] = [] + main_layers = self.hparams["num_hidden_layers"] + if not self.mtp_only: + consumed.extend(self._write_hash_routing_tensors()) + elif self.hparams["num_hash_layers"] > 0: + for bid in range(self.hparams["num_hash_layers"]): + name = f"layers.{bid}.ffn.gate.tid2eid" + if name in self.model_tensors: + consumed.extend(self._write_hash_routing_tensors()) + break + for bid in range(self.block_count): + if self.mtp_only and bid < main_layers: + continue consumed.extend(self._write_mxfp4_expert_tensor(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP)) consumed.extend(self._write_mxfp4_expert_tensor(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP)) consumed.extend(self._write_mxfp4_expert_tensor(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP)) + for bid in range(main_layers, self.block_count): + e_name = f"layers.{bid}.nextn.e_proj.weight" + h_name = f"layers.{bid}.nextn.h_proj.weight" + if e_name not in self.model_tensors and h_name not in self.model_tensors: + continue + if e_name not in self.model_tensors or h_name not in self.model_tensors: + raise KeyError(f"Missing DeepSeek-V4 MTP e/h projection pair for block {bid}") + + e_proj = LazyTorchTensor.to_eager(self.model_tensors[e_name]()) + h_proj = LazyTorchTensor.to_eager(self.model_tensors[h_name]()) + yield (f"layers.{bid}.nextn.eh_proj.weight", torch.cat((e_proj, h_proj), dim=1).contiguous()) + consumed.extend((e_name, h_name)) + for name in consumed: del self.model_tensors[name] @@ -737,6 +821,12 @@ def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_ "ffn.shared_experts.w1.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"), "ffn.shared_experts.w2.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"), "ffn.shared_experts.w3.weight": (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"), + "nextn.eh_proj.weight": (gguf.MODEL_TENSOR.NEXTN_EH_PROJ, ".weight"), + "nextn.enorm.weight": (gguf.MODEL_TENSOR.NEXTN_ENORM, ".weight"), + "nextn.hnorm.weight": (gguf.MODEL_TENSOR.NEXTN_HNORM, ".weight"), + "nextn.shared_head_norm.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ".weight"), + "nextn.embed_tokens.weight": (gguf.MODEL_TENSOR.NEXTN_EMBED_TOKENS, ".weight"), + "nextn.shared_head_head.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, ".weight"), } tensor_name = match.group(2) @@ -759,10 +849,12 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter return [(self._format_dsv4_tensor_name(tensor_key, bid, suffix), data_torch)] def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: - del new_name, bid # unused + del bid # unused if name in self._dsv4_fp8_dequantized and n_dims >= 2: return gguf.GGMLQuantizationType.Q8_0 + if new_name.endswith(".nextn.eh_proj.weight"): + return gguf.GGMLQuantizationType.Q8_0 if name in self._dsv4_f32_tensors: return gguf.GGMLQuantizationType.F32 if name in self._dsv4_bf16_tensors and n_dims >= 2: @@ -770,7 +862,122 @@ def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: return False + def prepare_metadata(self, vocab_only: bool): + from_dir = self.fname_out.is_dir() + super().prepare_metadata(vocab_only=vocab_only) + + if not self.mtp_only or not from_dir: + return + + output_type: str = self.ftype.name.partition("_")[2] + fname_default: str = gguf.naming_convention( + self.metadata.name, self.metadata.basename, self.metadata.finetune, + self.metadata.version, size_label=None, output_type=output_type, model_type=None) + self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" + def prepare_tensors(self): super().prepare_tensors() self._is_mxfp4 = True self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE + + +@ModelBase.register("DeepseekV4DSparkModel") +class DeepseekV4DSparkModel(DeepseekV4Model): + model_arch = gguf.MODEL_ARCH.DFLASH + + _DSPARK_ROOT_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = { + "main_proj.weight": (gguf.MODEL_TENSOR.FC, ".weight"), + "main_norm.weight": (gguf.MODEL_TENSOR.ENC_OUTPUT_NORM, ".weight"), + "markov_head.markov_w1.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W1, ".weight"), + "markov_head.markov_w2.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W2, ".weight"), + "confidence_head.proj.weight": (gguf.MODEL_TENSOR.DSPARK_CONF_PROJ, ".weight"), + } + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.block_count = 1 + max( + int(match.group(1)) for name in self.model_tensors + if (match := re.match(r"layers\.(\d+)\.", name)) + ) + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + self.hparams["compress_ratios"] = [0] * self.block_count + self.hparams["num_hash_layers"] = 0 + + def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]: + if remote_hf_model_id is None: + return super().index_tensors() + + with open(self.dir_model / "model.safetensors.index.json", "r", encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + + part_names = sorted({ + part_name for name, part_name in weight_map.items() + if name.startswith("mtp.") + }) + tensors: dict[str, Callable[[], Tensor]] = {} + + for part_name in part_names: + from huggingface_hub import hf_hub_download + + logger.info("gguf: caching remote DSpark part '%s'", part_name) + part_path = Path(hf_hub_download(repo_id=remote_hf_model_id, filename=part_name)) + with gguf.utility.SafetensorsLocal(part_path) as model_part: + for name in model_part: + data = model_part[name] + data_gen = lambda data=data: LazyTorchTensor.from_local_tensor(data) # noqa: E731 + if titem := self.filter_tensors((name, data_gen)): + tensor_name, tensor_gen = titem + tensors[tensor_name] = tensor_gen + + return tensors + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if not name.startswith("mtp."): + return None + return super().filter_tensors((cls._rekey_mtp_tensor_name(name), gen)) + + @staticmethod + def _rekey_mtp_tensor_name(name: str) -> str: + match = re.match(r"mtp\.(\d+)\.(.+)$", name) + if match is None: + raise ValueError(f"Unexpected DSpark tensor {name!r}") + + stage, rest = match.group(1), match.group(2) + root_names = ( + "main_proj.scale", + "norm.weight", + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + ) + if rest in DeepseekV4DSparkModel._DSPARK_ROOT_MAP or rest in root_names: + return rest + return f"layers.{stage}.{rest}" + + def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]: + if name in self._DSPARK_ROOT_MAP: + return self._DSPARK_ROOT_MAP[name] + return super()._map_dsv4_tensor_name(name, bid) + + def set_vocab(self): + if self.target_model_dir is None: + raise ValueError("DeepSeek-V4 DSpark requires --target-model-dir with the target tokenizer") + + original_dir = self.dir_model + try: + self.dir_model = self.target_model_dir + super().set_vocab() + finally: + self.dir_model = original_dir + + self.gguf_writer.add_mask_token_id(self.hparams["dspark_noise_token_id"]) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + self.gguf_writer.add_block_size(self.hparams["dspark_block_size"]) + self.gguf_writer.add_target_layers([layer + 1 for layer in self.hparams["dspark_target_layer_ids"]]) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 2c5e62a16fb..78ad26c6563 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -122,8 +122,12 @@ def parse_args() -> argparse.Namespace: help="Export only the multi-token prediction (MTP) head as a separate GGUF, suitable for use as a speculative draft. An 'mtp-' prefix will be added to the output file name.", ) parser.add_argument( - "--no-mtp", action="store_true", - help="Exclude the multi-token prediction (MTP) head from the converted GGUF. Pair with --mtp on a second run to publish trunk and MTP as two files. Note: the split form duplicates embeddings, but even though the bundled default is more space-efficient overall, this allows differing quantization which may be more performant.", + "--no-nextn", "--no-mtp", dest="no_mtp", action="store_true", + help="Exclude NextN speculative draft tensors from the converted GGUF. Pair with --mtp or --dspark on a second run to publish target and draft as two files.", + ) + parser.add_argument( + "--dspark", action="store_true", + help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.", ) parser.add_argument( "--mistral-format", action="store_true", @@ -254,13 +258,20 @@ def main() -> None: from conversion.mistral import MistralModel model_class = MistralModel - if args.mtp and args.no_mtp: - logger.error("--mtp and --no-mtp are mutually exclusive") + if sum((args.mtp, args.no_mtp, args.dspark)) > 1: + logger.error("--mtp, --no-nextn, and --dspark are mutually exclusive") sys.exit(1) + if args.dspark: + if is_mistral_format or model_architecture != "DeepseekV4ForCausalLM": + logger.error("--dspark is only supported for DeepseekV4ForCausalLM") + sys.exit(1) + from conversion.deepseek import DeepseekV4DSparkModel + model_class = DeepseekV4DSparkModel + if args.mtp or args.no_mtp: if not model_class.supports_mtp_export: - logger.error("--mtp / --no-mtp are not supported for %s", model_architecture) + logger.error("--mtp / --no-nextn are not supported for %s", model_architecture) sys.exit(1) if args.no_mtp: model_class.no_mtp = True diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 7fe2d669609..bfa5f134996 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -477,6 +477,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max(ggml_me return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, ggml_op op) { + const char * name = nullptr; + + switch (op) { + case GGML_OP_DSV4_HC_COMB: name = "kernel_dsv4_hc_comb_f32"; break; + case GGML_OP_DSV4_HC_PRE: name = "kernel_dsv4_hc_pre_f32"; break; + case GGML_OP_DSV4_HC_POST: name = "kernel_dsv4_hc_post_f32"; break; + default: GGML_ABORT("fatal error"); + } + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv(ggml_metal_library_t lib, const ggml_tensor * op) { GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); @@ -2260,6 +2278,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_opt_step_sgd(ggm return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back(ggml_metal_library_t lib, const ggml_tensor * op) { + assert(op->op == GGML_OP_SILU_BACK); + + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_silu_back_%s", ggml_type_name(op->src[0]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_memset(ggml_metal_library_t lib, const ggml_tensor * op) { GGML_ASSERT(op->type == GGML_TYPE_I64); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index ee2116a0628..e21fe5ab67d 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -117,6 +117,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat (ggml_metal_library_t lib, enum ggml_type tsrc); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat (ggml_metal_library_t lib, enum ggml_type tsrc); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows (ggml_metal_library_t lib, const struct ggml_tensor * op); @@ -124,6 +125,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_bl struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, enum ggml_op op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 4db05219a6b..eb7a74aa60d 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1174,6 +1174,14 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te default: return false; } + case GGML_OP_SILU_BACK: + return (op->src[0]->type == GGML_TYPE_F32) && + (op->src[1]->type == GGML_TYPE_F32) && + (op->type == GGML_TYPE_F32) && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]) && + ggml_is_contiguous(op) && + ggml_are_same_shape(op->src[0], op->src[1]); case GGML_OP_GLU: switch (ggml_get_glu_op(op)) { case GGML_GLU_OP_REGLU: @@ -1218,6 +1226,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_MUL: case GGML_OP_DIV: case GGML_OP_ADD_ID: + return ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]) && (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) && (op->src[0]->type == op->src[1]->type); case GGML_OP_ACC: return ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]) && op->src[0]->type == GGML_TYPE_F32; case GGML_OP_REPEAT: @@ -1372,6 +1381,42 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return false; } return has_simdgroup_mm; // TODO: over-restricted for vec-kernels + case GGML_OP_DSV4_HC_COMB: + return has_simdgroup_reduction && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == 24 && + op->src[1]->ne[0] >= 3 && + op->src[2]->ne[0] == 24 && + ggml_is_contiguous_rows(op->src[0]) && + ggml_is_contiguous_rows(op->src[1]) && + ggml_is_contiguous_rows(op->src[2]); + case GGML_OP_DSV4_HC_PRE: + return has_simdgroup_reduction && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + op->src[0]->ne[1] == 4 && + op->src[1]->ne[0] == 4 && + ggml_is_contiguous_rows(op->src[0]) && + ggml_is_contiguous_rows(op->src[1]); + case GGML_OP_DSV4_HC_POST: + return has_simdgroup_reduction && + op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && + op->src[3]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + op->src[1]->ne[1] == 4 && + op->src[2]->ne[0] == 4 && + op->src[3]->ne[0] == 4 && + op->src[3]->ne[1] == 4 && + ggml_is_contiguous_rows(op->src[0]) && + ggml_is_contiguous_rows(op->src[1]) && + ggml_is_contiguous_rows(op->src[2]) && + ggml_is_contiguous_rows(op->src[3]); case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: return has_simdgroup_reduction; diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 9a214a708a8..c29f7a80eaf 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -1219,6 +1219,49 @@ typedef struct { int64_t val; } ggml_metal_kargs_memset; +typedef struct { + int32_t n_tokens; + int32_t n_iter; + uint64_t nb_m0; + uint64_t nb_m1; + uint64_t nb_s0; + uint64_t nb_b0; + uint64_t nb_d0; + uint64_t nb_d1; + uint64_t nb_d2; + float eps; +} ggml_metal_kargs_dsv4_hc_comb; + +typedef struct { + int32_t n_embd; + int32_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_x2; + uint64_t nb_w0; + uint64_t nb_w1; + uint64_t nb_d0; + uint64_t nb_d1; +} ggml_metal_kargs_dsv4_hc_pre; + +typedef struct { + int32_t n_embd; + int32_t n_tokens; + uint64_t nb_x0; + uint64_t nb_x1; + uint64_t nb_r0; + uint64_t nb_r1; + uint64_t nb_r2; + uint64_t nb_p0; + uint64_t nb_p1; + uint64_t nb_c0; + uint64_t nb_c1; + uint64_t nb_c2; + uint64_t nb_d0; + uint64_t nb_d1; + uint64_t nb_d2; +} ggml_metal_kargs_dsv4_hc_post; + typedef struct { int32_t ne00; int32_t ne01; @@ -1270,4 +1313,8 @@ typedef struct { int64_t np; } ggml_metal_kargs_opt_step_sgd; +typedef struct { + int64_t ne; +} ggml_metal_kargs_silu_back; + #endif // GGML_METAL_IMPL diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 55915ae87c8..043281e21ac 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -328,6 +328,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_unary(ctx, idx); } break; + case GGML_OP_SILU_BACK: + { + n_fuse = ggml_metal_op_silu_back(ctx, idx); + } break; case GGML_OP_GLU: { n_fuse = ggml_metal_op_glu(ctx, idx); @@ -345,6 +349,12 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_cumsum(ctx, idx); } break; + case GGML_OP_DSV4_HC_COMB: + case GGML_OP_DSV4_HC_PRE: + case GGML_OP_DSV4_HC_POST: + { + n_fuse = ggml_metal_op_dsv4_hc(ctx, idx); + } break; case GGML_OP_SOFT_MAX: { n_fuse = ggml_metal_op_soft_max(ctx, idx); @@ -1330,6 +1340,137 @@ int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_encoder_t enc = ctx->enc; + auto pipeline = ggml_metal_library_get_pipeline_dsv4_hc(ctx->lib, op->op); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + + switch (op->op) { + case GGML_OP_DSV4_HC_COMB: + { + const ggml_tensor * mixes = op->src[0]; + const ggml_tensor * scale = op->src[1]; + const ggml_tensor * base = op->src[2]; + + GGML_ASSERT(mixes->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + GGML_ASSERT(mixes->ne[0] == 24); + GGML_ASSERT(op->ne[0] == 4 && op->ne[1] == 4); + + ggml_metal_kargs_dsv4_hc_comb args = { + /*.n_tokens =*/ (int32_t) mixes->ne[1], + /*.n_iter =*/ ggml_get_op_params_i32(op, 1), + /*.nb_m0 =*/ mixes->nb[0], + /*.nb_m1 =*/ mixes->nb[1], + /*.nb_s0 =*/ scale->nb[0], + /*.nb_b0 =*/ base->nb[0], + /*.nb_d0 =*/ op->nb[0], + /*.nb_d1 =*/ op->nb[1], + /*.nb_d2 =*/ op->nb[2], + /*.eps =*/ ggml_get_op_params_f32(op, 0), + }; + + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(mixes), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(scale), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(base), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); + + // One SIMDgroup owns one 4x4 Sinkhorn matrix. Packing up to four + // independent tokens per threadgroup keeps both decode and prompt + // dispatches compact without any threadgroup-memory synchronization. + const int nsg = std::min(4, args.n_tokens); + ggml_metal_encoder_dispatch_threadgroups( + enc, (args.n_tokens + nsg - 1)/nsg, 1, 1, 32, nsg, 1); + } break; + case GGML_OP_DSV4_HC_PRE: + { + const ggml_tensor * x = op->src[0]; + const ggml_tensor * weights = op->src[1]; + + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + GGML_ASSERT(x->ne[1] == 4); + + ggml_metal_kargs_dsv4_hc_pre args = { + /*.n_embd =*/ (int32_t) x->ne[0], + /*.n_tokens =*/ (int32_t) x->ne[2], + /*.nb_x0 =*/ x->nb[0], + /*.nb_x1 =*/ x->nb[1], + /*.nb_x2 =*/ x->nb[2], + /*.nb_w0 =*/ weights->nb[0], + /*.nb_w1 =*/ weights->nb[1], + /*.nb_d0 =*/ op->nb[0], + /*.nb_d1 =*/ op->nb[1], + }; + + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weights), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3); + + const int n_tiles = (args.n_embd + 31)/32; + const int nsg = std::min(4, n_tiles); + ggml_metal_encoder_dispatch_threadgroups( + enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1); + } break; + case GGML_OP_DSV4_HC_POST: + { + const ggml_tensor * x = op->src[0]; + const ggml_tensor * residual = op->src[1]; + const ggml_tensor * post = op->src[2]; + const ggml_tensor * comb = op->src[3]; + + GGML_ASSERT(x->type == GGML_TYPE_F32); + GGML_ASSERT(residual->type == GGML_TYPE_F32); + GGML_ASSERT(post->type == GGML_TYPE_F32); + GGML_ASSERT(comb->type == GGML_TYPE_F32); + GGML_ASSERT(op->type == GGML_TYPE_F32); + GGML_ASSERT(residual->ne[1] == 4); + + ggml_metal_kargs_dsv4_hc_post args = { + /*.n_embd =*/ (int32_t) x->ne[0], + /*.n_tokens =*/ (int32_t) x->ne[1], + /*.nb_x0 =*/ x->nb[0], + /*.nb_x1 =*/ x->nb[1], + /*.nb_r0 =*/ residual->nb[0], + /*.nb_r1 =*/ residual->nb[1], + /*.nb_r2 =*/ residual->nb[2], + /*.nb_p0 =*/ post->nb[0], + /*.nb_p1 =*/ post->nb[1], + /*.nb_c0 =*/ comb->nb[0], + /*.nb_c1 =*/ comb->nb[1], + /*.nb_c2 =*/ comb->nb[2], + /*.nb_d0 =*/ op->nb[0], + /*.nb_d1 =*/ op->nb[1], + /*.nb_d2 =*/ op->nb[2], + }; + + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(residual), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(post), 3); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comb), 4); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 5); + + const int n_tiles = (args.n_embd + 31)/32; + const int nsg = std::min(4, n_tiles); + ggml_metal_encoder_dispatch_threadgroups( + enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1); + } break; + default: + GGML_ABORT("fatal error"); + } + + return 1; +} + int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -3615,9 +3756,6 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { GGML_TENSOR_LOCALS( int32_t, ne, op, ne); GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); - GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); - GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); GGML_ASSERT(ggml_is_contiguous_rows(op->src[1])); @@ -3757,6 +3895,36 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { return n_fuse; } +int ggml_metal_op_silu_back(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + auto pipeline = ggml_metal_library_get_pipeline_silu_back(lib, op); + + const int64_t ne = ggml_nelements(op); + + ggml_metal_kargs_silu_back args = { + /*.ne =*/ ne, + }; + + int arg_idx{0}; + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), arg_idx++); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), arg_idx++); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), arg_idx++); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), arg_idx++); + + const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne); + const int64_t n = (ne + nth - 1) / nth; + + ggml_metal_encoder_dispatch_threadgroups(enc, n, 1, 1, nth, 1, 1); + + return 1; +} + int ggml_metal_op_l2_norm(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index e7a343d8bfd..0abb634a12c 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -54,6 +54,7 @@ int ggml_metal_op_cumsum (ggml_metal_op_t ctx, int idx); int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx); int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx); int ggml_metal_op_soft_max (ggml_metal_op_t ctx, int idx); int ggml_metal_op_ssm_conv (ggml_metal_op_t ctx, int idx); int ggml_metal_op_ssm_scan (ggml_metal_op_t ctx, int idx); @@ -71,6 +72,7 @@ int ggml_metal_op_mul_mat_id (ggml_metal_op_t ctx, int idx); int ggml_metal_op_add_id (ggml_metal_op_t ctx, int idx); int ggml_metal_op_flash_attn_ext (ggml_metal_op_t ctx, int idx); int ggml_metal_op_bin (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_silu_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_l2_norm (ggml_metal_op_t ctx, int idx); int ggml_metal_op_group_norm (ggml_metal_op_t ctx, int idx); int ggml_metal_op_norm (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 54429a6b8a5..6431cd132ed 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -2046,6 +2046,20 @@ template [[host_name("kernel_unary_f32_f32_4")]] kernel kernel_unary_t kernel_un template [[host_name("kernel_unary_f16_f16")]] kernel kernel_unary_t kernel_unary_impl; template [[host_name("kernel_unary_f16_f16_4")]] kernel kernel_unary_t kernel_unary_impl; +kernel void kernel_silu_back_f32( + constant ggml_metal_kargs_silu_back & args, + device const float * dy, + device const float * x, + device float * dx, + uint gid [[thread_position_in_grid]]) { + if (gid >= args.ne) { + return; + } + + const float s = 1.0f / (1.0f + exp(-x[gid])); + dx[gid] = dy[gid] * s * (1.0f + x[gid] * (1.0f - s)); +} + // OP: 0 - add, 1 - sub, 2 - mul, 3 - div constant short FC_bin_op [[function_constant(FC_BIN + 0)]]; constant short FC_bin_f [[function_constant(FC_BIN + 1)]]; @@ -2209,6 +2223,8 @@ typedef decltype(kernel_bin_fuse_impl) kernel_bin_fuse_t; template [[host_name("kernel_bin_fuse_f32_f32_f32")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl; template [[host_name("kernel_bin_fuse_f32_f32_f32_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl; +template [[host_name("kernel_bin_fuse_f16_f16_f16")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl; +template [[host_name("kernel_bin_fuse_f16_f16_f16_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl; kernel void kernel_add_id( constant ggml_metal_kargs_add_id & args, @@ -13465,3 +13481,162 @@ kernel void kernel_count_equal( typedef decltype(kernel_count_equal) kernel_count_equal_t; template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal; + +kernel void kernel_dsv4_hc_comb_f32( + constant ggml_metal_kargs_dsv4_hc_comb & args, + device const char * mixes, + device const char * scale, + device const char * base, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort hc = 4; + constexpr ushort comb_offset = 2*hc; + + const int it = tgpig.x*ntg.y + sgitg; + if (it >= args.n_tokens) { + return; + } + + float scale_lane = 0.0f; + if (tiisg == 0) { + scale_lane = *(device const float *) (scale + 2*args.nb_s0); + } + const float scale_comb = simd_shuffle(scale_lane, 0); + + float v = 0.0f; + if (tiisg < hc*hc) { + v = *(device const float *) (mixes + (comb_offset + tiisg)*args.nb_m0 + it*args.nb_m1)*scale_comb + + *(device const float *) (base + (comb_offset + tiisg)*args.nb_b0); + } + + // Softmax across destinations (the four contiguous lanes for each source). + float vmax = max(v, simd_shuffle_xor(v, 1)); + vmax = max(vmax, simd_shuffle_xor(vmax, 2)); + v = exp(v - vmax); + + float sum = v + simd_shuffle_xor(v, 1); + sum += simd_shuffle_xor(sum, 2); + v = v/sum + args.eps; + + // Normalize columns: equal destination indices are four lanes apart. + sum = v + simd_shuffle_xor(v, 4); + sum += simd_shuffle_xor(sum, 8); + v /= sum + args.eps; + + for (int i = 1; i < args.n_iter; ++i) { + sum = v + simd_shuffle_xor(v, 1); + sum += simd_shuffle_xor(sum, 2); + v /= sum + args.eps; + + sum = v + simd_shuffle_xor(v, 4); + sum += simd_shuffle_xor(sum, 8); + v /= sum + args.eps; + } + + if (tiisg < hc*hc) { + const ushort idst = tiisg & 3; + const ushort isrc = tiisg >> 2; + *(device float *) (dst + idst*args.nb_d0 + isrc*args.nb_d1 + it*args.nb_d2) = v; + } +} + +kernel void kernel_dsv4_hc_pre_f32( + constant ggml_metal_kargs_dsv4_hc_pre & args, + device const char * x, + device const char * weights, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort hc = 4; + + const int it = tgpig.y; + const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg; + + float weight_lane = 0.0f; + if (tiisg < hc) { + weight_lane = *(device const float *) (weights + tiisg*args.nb_w0 + it*args.nb_w1); + } + + float w[hc]; + FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) { + w[ih] = simd_shuffle(weight_lane, ih); + } + + if (i0 >= args.n_embd) { + return; + } + + device const char * xb = x + i0*args.nb_x0 + it*args.nb_x2; + float result = 0.0f; + FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) { + result = fma(*(device const float *) (xb + ih*args.nb_x1), w[ih], result); + } + + *(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = result; +} + +kernel void kernel_dsv4_hc_post_f32( + constant ggml_metal_kargs_dsv4_hc_post & args, + device const char * x, + device const char * residual, + device const char * post, + device const char * comb, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + constexpr ushort hc = 4; + + const int it = tgpig.y; + const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg; + + float coeff_lane = 0.0f; + if (tiisg < hc) { + coeff_lane = *(device const float *) (post + tiisg*args.nb_p0 + it*args.nb_p1); + } else if (tiisg < hc + hc*hc) { + const ushort idx = tiisg - hc; + const ushort idst = idx & 3; + const ushort isrc = idx >> 2; + coeff_lane = *(device const float *) (comb + idst*args.nb_c0 + isrc*args.nb_c1 + it*args.nb_c2); + } + + float post_reg[hc]; + float comb_reg[hc][hc]; + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + post_reg[idst] = simd_shuffle(coeff_lane, idst); + } + FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) { + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + comb_reg[isrc][idst] = simd_shuffle(coeff_lane, hc + idst + hc*isrc); + } + } + + if (i0 >= args.n_embd) { + return; + } + + const float xv = *(device const float *) (x + i0*args.nb_x0 + it*args.nb_x1); + float result[hc]; + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + result[idst] = xv*post_reg[idst]; + } + + device const char * rb = residual + i0*args.nb_r0 + it*args.nb_r2; + FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) { + const float rv = *(device const float *) (rb + isrc*args.nb_r1); + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + result[idst] = fma(rv, comb_reg[isrc][idst], result[idst]); + } + } + + FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) { + *(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst]; + } +} + diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index d07b8fe41a3..915c8e7b903 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -7495,6 +7495,7 @@ static ggml_backend_i ggml_backend_opencl_i = { ggml_backend_t ggml_backend_opencl_init(void) { ggml_backend_dev_t dev = ggml_backend_reg_dev_get(ggml_backend_opencl_reg(), 0); ggml_backend_opencl_context *backend_ctx = ggml_cl_init(dev); + backend_ctx->ref_count++; ggml_backend_t backend = new ggml_backend { /* .guid = */ ggml_backend_opencl_guid(), @@ -24259,7 +24260,7 @@ static void ggml_cl_glu(ggml_backend_t backend, const ggml_tensor * src0, const } const size_t nrows = ggml_nrows(src0); - size_t nth = 512; + size_t nth = backend_ctx->max_workgroup_size < 512 ? backend_ctx->max_workgroup_size : 512; size_t global_work_size[] = {nrows*nth, 1, 1}; size_t local_work_size[] = {nth, 1, 1}; diff --git a/ggml/src/ggml-sycl/common.hpp b/ggml/src/ggml-sycl/common.hpp index c57412a67d8..8d338c6ed68 100644 --- a/ggml/src/ggml-sycl/common.hpp +++ b/ggml/src/ggml-sycl/common.hpp @@ -235,6 +235,7 @@ struct sycl_device_info { int max_wg_per_cu; // max work groups per compute unit - refer to // cudaOccupancyMaxActiveBlocksPerMultiprocessor bool vmm; // virtual memory support + bool l0_device_type_valid; bool l0_discrete_gpu; // Level Zero backend and not an integrated GPU size_t vmm_granularity; // granularity of virtual memory size_t total_vram; diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 3b7ba09c48c..131bff714f9 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -168,7 +168,10 @@ static ggml_sycl_device_info ggml_sycl_init() { ze_device_properties_t props = {}; props.stype = ZE_STRUCTURE_TYPE_DEVICE_PROPERTIES; ze_result_t r = zeDeviceGetProperties(ze_dev, &props); - info.devices[i].l0_discrete_gpu = r == ZE_RESULT_SUCCESS && !(props.flags & ZE_DEVICE_PROPERTY_FLAG_INTEGRATED); + if (r == ZE_RESULT_SUCCESS) { + info.devices[i].l0_device_type_valid = true; + info.devices[i].l0_discrete_gpu = !(props.flags & ZE_DEVICE_PROPERTY_FLAG_INTEGRATED); + } } #endif } @@ -5610,7 +5613,11 @@ static void ggml_backend_sycl_device_get_memory(ggml_backend_dev_t dev, size_t * } static enum ggml_backend_dev_type ggml_backend_sycl_device_get_type(ggml_backend_dev_t dev) { - GGML_UNUSED(dev); + ggml_backend_sycl_device_context * ctx = (ggml_backend_sycl_device_context *)dev->context; + const sycl_device_info & info = ggml_sycl_info().devices[ctx->device]; + if (info.l0_device_type_valid && !info.l0_discrete_gpu) { + return GGML_BACKEND_DEVICE_TYPE_IGPU; + } return GGML_BACKEND_DEVICE_TYPE_GPU; } diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 1a08e8d6476..af0d4629f86 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -610,6 +610,13 @@ static constexpr std::initializer_list topk_moe_sigmoid_norm_bias{ GGML GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE }; +static constexpr std::initializer_list topk_moe_sqrt_softplus_norm_bias{ GGML_OP_UNARY, GGML_OP_SQRT, + GGML_OP_RESHAPE, GGML_OP_ADD, + GGML_OP_ARGSORT, GGML_OP_VIEW, + GGML_OP_GET_ROWS, GGML_OP_RESHAPE, + GGML_OP_SUM_ROWS, GGML_OP_CLAMP, + GGML_OP_DIV, GGML_OP_RESHAPE }; + static constexpr std::initializer_list topk_moe_early_softmax { GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }; @@ -673,6 +680,22 @@ static constexpr std::initializer_list> topk_moe_sigmoid_norm {10, 0, 9 }, // reshape->src[0] == div }; +static constexpr std::initializer_list> topk_moe_sqrt_softplus_norm_bias_edges { + { 1, 0, 0 }, // sqrt->src[0] == softplus + { 2, 0, 1 }, // reshape->src[0] == sqrt + { 3, 0, 1 }, // add->src[0] == sqrt + { 4, 0, 3 }, // argsort->src[0] == add + { 5, 0, 4 }, // view->src[0] == argsort + { 6, 0, 2 }, // get_rows->src[0] == reshape + { 6, 1, 5 }, // get_rows->src[1] == view + { 7, 0, 6 }, // reshape->src[0] == get_rows + { 8, 0, 7 }, // sum_rows->src[0] == reshape + { 9, 0, 8 }, // clamp->src[0] == sum_rows + {10, 0, 7 }, // div->src[0] == reshape + {10, 1, 9 }, // div->src[1] == clamp + {11, 0,10 }, // reshape->src[0] == div +}; + // same as early_softmax_norm but ending after the get_rows static constexpr std::initializer_list> topk_moe_early_softmax_edges { { 1, 0, 0 }, // reshape->src[0] == softmax @@ -701,6 +724,7 @@ enum topk_moe_mode { TOPK_MOE_EARLY_SOFTMAX_NORM, TOPK_MOE_LATE_SOFTMAX, TOPK_MOE_SIGMOID_NORM_BIAS, + TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS, TOPK_MOE_COUNT, }; @@ -13248,12 +13272,16 @@ static void ggml_vk_soft_max_back(ggml_backend_vk_context * ctx, vk_context& sub static void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { topk_moe_mode mode = ctx->fused_topk_moe_mode; + const bool has_bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS || mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS; ggml_tensor * logits = cgraph->nodes[node_idx + 0]->src[0]; - ggml_tensor * bias = (mode == TOPK_MOE_SIGMOID_NORM_BIAS) ? cgraph->nodes[node_idx + 2]->src[1] : logits; + ggml_tensor * bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? cgraph->nodes[node_idx + 2]->src[1] : + mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? cgraph->nodes[node_idx + 3]->src[1] : + logits; ggml_tensor * weights = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; - ggml_tensor * ids = (mode == TOPK_MOE_SIGMOID_NORM_BIAS) ? cgraph->nodes[node_idx + 4] : - (mode == TOPK_MOE_LATE_SOFTMAX) ? cgraph->nodes[node_idx + 1] : - cgraph->nodes[node_idx + 3]; + ggml_tensor * ids = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? cgraph->nodes[node_idx + 4] : + mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? cgraph->nodes[node_idx + 5] : + mode == TOPK_MOE_LATE_SOFTMAX ? cgraph->nodes[node_idx + 1] : + cgraph->nodes[node_idx + 3]; GGML_ASSERT(logits->type == GGML_TYPE_F32); GGML_ASSERT(bias->type == GGML_TYPE_F32); @@ -13293,16 +13321,24 @@ static void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx, pc.clamp_min = ggml_get_op_params_f32(clamp, 0); pc.clamp_max = ggml_get_op_params_f32(clamp, 1); } + if (mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS) { + ggml_tensor * clamp = cgraph->nodes[node_idx + 9]; + GGML_ASSERT(clamp->op == GGML_OP_CLAMP); + pc.clamp_min = ggml_get_op_params_f32(clamp, 0); + pc.clamp_max = ggml_get_op_params_f32(clamp, 1); + } #define GATING_FUNC_SOFTMAX 0 #define GATING_FUNC_SIGMOID 1 #define GATING_FUNC_SOFTMAX_WEIGHT 2 - - pc.gating_func = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? GATING_FUNC_SIGMOID : - mode == TOPK_MOE_LATE_SOFTMAX ? GATING_FUNC_SOFTMAX_WEIGHT : - GATING_FUNC_SOFTMAX; - pc.has_bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS; - pc.with_norm = mode == TOPK_MOE_EARLY_SOFTMAX_NORM || mode == TOPK_MOE_SIGMOID_NORM_BIAS; +#define GATING_FUNC_SQRT_SOFTPLUS 3 + + pc.gating_func = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? GATING_FUNC_SIGMOID : + mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? GATING_FUNC_SQRT_SOFTPLUS : + mode == TOPK_MOE_LATE_SOFTMAX ? GATING_FUNC_SOFTMAX_WEIGHT : + GATING_FUNC_SOFTMAX; + pc.has_bias = has_bias; + pc.with_norm = mode == TOPK_MOE_EARLY_SOFTMAX_NORM || has_bias; if (ctx->fused_topk_moe_scale) { GGML_ASSERT(weights->op == GGML_OP_SCALE); pc.output_scale = ggml_get_op_params_f32(weights, 0); @@ -16415,6 +16451,20 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc return false; } break; + case TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS: + softmax = cgraph->nodes[node_idx + 0]; // really softplus + weights = cgraph->nodes[node_idx + 11]; + get_rows = cgraph->nodes[node_idx + 6]; + argsort = cgraph->nodes[node_idx + 4]; + if (ggml_get_unary_op(softmax) != GGML_UNARY_OP_SOFTPLUS) { + return false; + } + // bias is expected to be 1D + if (ggml_nrows(cgraph->nodes[node_idx + 3]->src[1]) != 1 || + !ggml_is_contiguous(cgraph->nodes[node_idx + 3]->src[1])) { + return false; + } + break; case TOPK_MOE_EARLY_SOFTMAX: softmax = cgraph->nodes[node_idx + 0]; weights = cgraph->nodes[node_idx + 4]; @@ -16438,7 +16488,9 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc probs = probs->src[0]; ggml_tensor * selection_probs = argsort->src[0]; - if (probs != selection_probs && mode != TOPK_MOE_SIGMOID_NORM_BIAS) { + if (probs != selection_probs && + mode != TOPK_MOE_SIGMOID_NORM_BIAS && + mode != TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS) { return false; } @@ -16806,7 +16858,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg // the fused result in an elementwise-way. This affects whether the memory for // the src is allowed to overlap the memory for the destination. // The array is sized to handle the largest fusion (asserted later). - bool op_srcs_fused_elementwise[12]; + bool op_srcs_fused_elementwise[13]; ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; @@ -16917,6 +16969,15 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_SIGMOID_NORM_BIAS; fusion_string = "TOPK_MOE_SIGMOID_NORM_BIAS"; std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false); + } else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_sqrt_softplus_norm_bias, { i + 5, i + 11 }) && + ggml_check_edges(cgraph, i, topk_moe_sqrt_softplus_norm_bias_edges) && + ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS)) { + ctx->num_additional_fused_ops = topk_moe_sqrt_softplus_norm_bias.size() - 1; + // view of argsort writes to memory + ctx->fused_ops_write_mask |= 1 << 5; + ctx->fused_topk_moe_mode = TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS; + fusion_string = "TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS"; + std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false); } else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_early_softmax, { i + 3, i + 4 }) && ggml_check_edges(cgraph, i, topk_moe_early_softmax_edges) && ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_EARLY_SOFTMAX)) { @@ -17183,6 +17244,9 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * if (keep_pattern(topk_moe_sigmoid_norm_bias)) { continue; } + if (keep_pattern(topk_moe_sqrt_softplus_norm_bias)) { + continue; + } if (keep_pattern(topk_moe_early_softmax)) { continue; } @@ -17213,6 +17277,7 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * // Don't pull forward nodes from fusion patterns if (match_pattern(topk_moe_early_softmax_norm, j) || match_pattern(topk_moe_sigmoid_norm_bias, j) || + match_pattern(topk_moe_sqrt_softplus_norm_bias, j) || match_pattern(topk_moe_early_softmax, j) || match_pattern(topk_moe_late_softmax, j) || match_pattern(snake_pattern, j)) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/topk_moe.comp b/ggml/src/ggml-vulkan/vulkan-shaders/topk_moe.comp index ef2f202ec9b..d219201fda0 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/topk_moe.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/topk_moe.comp @@ -10,6 +10,7 @@ #define GATING_FUNC_SOFTMAX 0 #define GATING_FUNC_SIGMOID 1 #define GATING_FUNC_SOFTMAX_WEIGHT 2 +#define GATING_FUNC_SQRT_SOFTPLUS 3 layout (push_constant) uniform parameter { @@ -120,6 +121,13 @@ void main() { const uint expert = i + lane; probs[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? 1.f / (1.f + exp(-probs[i / WARP_SIZE])) : -INFINITY; } + } else if (gating_func == GATING_FUNC_SQRT_SOFTPLUS) { + [[unroll]] + for (uint i = 0; i < n_experts; i += WARP_SIZE) { + const uint expert = i + lane; + const float val = probs[i / WARP_SIZE]; + probs[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? sqrt(val > 20.0f ? val : log(1.0f + exp(val))) : -INFINITY; + } } float selection_probs[experts_per_thread]; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index babaddb6542..66c1c3c8977 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -2774,6 +2774,10 @@ class ggml_webgpu_shader_lib { defines.push_back("TYPE_F32"); variant += "_f32"; break; + case GGML_TYPE_F16: + defines.push_back("TYPE_F16"); + variant += "_f16"; + break; case GGML_TYPE_I32: defines.push_back("TYPE_I32"); variant += "_i32"; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index 370f05dfe67..c001cda7d11 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4290,7 +4290,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const supports_op = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_I32); break; case GGML_OP_REPEAT: - supports_op = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_I32 || src0->type == GGML_TYPE_I16); + supports_op = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_I32 || + src0->type == GGML_TYPE_I16); break; case GGML_OP_CPY: case GGML_OP_CONT: diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/repeat.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/repeat.wgsl index 6e2a1a8b614..43b883e6794 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/repeat.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/repeat.wgsl @@ -27,6 +27,9 @@ struct Params { #ifdef TYPE_I32 #define DataType i32 #endif +#ifdef TYPE_F16 +#define DataType f16 +#endif #ifdef TYPE_I16 // same size (16-bit) is sufficient for repeat #define DataType f16 diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f50d8b8f6eb..769de476c62 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -3337,6 +3337,12 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE_SHEXP, MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], MODEL_ARCH.ERNIE4_5_MOE: [ MODEL_TENSOR.TOKEN_EMBD, @@ -4384,10 +4390,37 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K_NORM, MODEL_TENSOR.ATTN_GATE, + + 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.HC_HEAD_FN, + MODEL_TENSOR.HC_HEAD_BASE, + MODEL_TENSOR.HC_HEAD_SCALE, + MODEL_TENSOR.FFN_NORM, MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + 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.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, # optional DSpark heads diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 4a3e52afff2..3320975d64f 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1017,6 +1017,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_DEEPSEEK4: return true; default: return false; @@ -1039,6 +1040,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { switch (arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_DEEPSEEK4: return true; default: return false; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index fe8ba7aa53d..386c8ee49eb 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -121,8 +121,9 @@ llama_context::llama_context( cparams.no_perf = params.no_perf; cparams.warmup = false; - cparams.embeddings_layer_inp.resize(hparams.n_layer(), false); - embd_layer_inp.resize(hparams.n_layer()); + // +1: id n_layer() taps the output of the last layer ("input" of the head) + cparams.embeddings_layer_inp.resize(hparams.n_layer() + 1, false); + embd_layer_inp.resize(hparams.n_layer() + 1); cparams.ctx_type = params.ctx_type; cparams.pooling_type = params.pooling_type; @@ -1219,7 +1220,7 @@ void llama_context::set_embeddings_nextn(bool value, bool masked) { void llama_context::set_embeddings_layer_inp(uint32_t lid, bool enable) { LLAMA_LOG_DEBUG("%s: lid = %d, enable = %d\n", __func__, lid, enable); - GGML_ASSERT(lid < model.hparams.n_layer()); + GGML_ASSERT(lid <= model.hparams.n_layer()); cparams.embeddings_layer_inp[lid] = enable; @@ -1771,7 +1772,8 @@ int llama_context::decode(const llama_batch & batch_inp) { const auto & hparams = model.hparams; const int64_t n_vocab = vocab.n_tokens(); - const int64_t n_embd = hparams.n_embd_inp(); + const bool mtp_embd = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && batch_inp.embd; + const int64_t n_embd = mtp_embd ? hparams.n_embd_out() : hparams.n_embd_inp(); // when computing embeddings, all tokens are output const bool output_all = cparams.embeddings; @@ -2328,8 +2330,9 @@ void llama_context::extract_layer_inputs(const llm_graph_result * res, size_t to } void llama_context::output_reorder() { - const uint64_t n_vocab = model.vocab.n_tokens(); - const uint64_t n_embd = model.hparams.n_embd; + const uint64_t n_vocab = model.vocab.n_tokens(); + const uint64_t n_embd = model.hparams.n_embd; + const uint64_t n_embd_out = model.hparams.n_embd_out(); for (size_t s = 0; s < output_swaps.size(); ++s) { const uint64_t i0 = output_swaps[s].i0; @@ -2342,14 +2345,14 @@ void llama_context::output_reorder() { } if (embd.size > 0) { - for (uint64_t k = 0; k < n_embd; k++) { - std::swap(embd.data[i0*n_embd + k], embd.data[i1*n_embd + k]); + for (uint64_t k = 0; k < n_embd_out; k++) { + std::swap(embd.data[i0*n_embd_out + k], embd.data[i1*n_embd_out + k]); } } if (embd_nextn.size > 0) { - for (uint64_t k = 0; k < n_embd; k++) { - std::swap(embd_nextn.data[i0*n_embd + k], embd_nextn.data[i1*n_embd + k]); + for (uint64_t k = 0; k < n_embd_out; k++) { + std::swap(embd_nextn.data[i0*n_embd_out + k], embd_nextn.data[i1*n_embd_out + k]); } } @@ -2399,7 +2402,15 @@ void llama_context::output_reorder() { // uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { - if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4) { + if (model.arch == LLM_ARCH_QWEN3NEXT || + model.arch == LLM_ARCH_KIMI_LINEAR || + model.arch == LLM_ARCH_QWEN35 || + model.arch == LLM_ARCH_QWEN35MOE || + model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_DFLASH || + model.arch == LLM_ARCH_NANBEIGE || + model.arch == LLM_ARCH_MINIMAX_M3) { + return std::max(n_tokens * 40, 32u * model.n_tensors()); } uint32_t res = std::max(1024u, 8u*model.n_tensors()); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8d33460c9d7..dff5ab68db5 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -585,20 +585,154 @@ bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) { // dsv4 helpers + +void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) { + // base tensors may not be allocated if there are no non-SWA attention layers + if (self_k_idxs && self_k_idxs->buffer) { + mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch); + if (self_v_idxs) { + mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch); + } + } + + // the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live + if (self_kq_mask && self_kq_mask->buffer) { + mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + } + + // swa tensors may not be allocated if there are no SWA attention layers + if (self_k_idxs_swa && self_k_idxs_swa->buffer) { + mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch); + if (self_v_idxs_swa) { + mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch); + } + } + + if (self_kq_mask_swa && self_kq_mask_swa->buffer) { + mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); + } + + if (self_k_rot && self_k_rot->buffer) { + mctx->get_base()->set_input_k_rot(self_k_rot); + } + + if (self_v_rot && self_v_rot->buffer) { + mctx->get_base()->set_input_v_rot(self_v_rot); + } + + if (self_k_rot_swa && self_k_rot_swa->buffer) { + mctx->get_swa()->set_input_k_rot(self_k_rot_swa); + } + + if (self_v_rot_swa && self_v_rot_swa->buffer) { + mctx->get_swa()->set_input_v_rot(self_v_rot_swa); + } +} + +bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { + const auto * mctx = static_cast(params.mctx); + + this->mctx = mctx; + + bool res = true; + + // base tensors may not be allocated if there are no non-SWA attention layers + if (self_k_idxs && self_k_idxs->buffer) { + res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; + //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there + } + + if (self_kq_mask && self_kq_mask->buffer) { + res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams); + } + + // swa tensors may not be allocated if there are no SWA attention layers + if (self_k_idxs_swa && self_k_idxs_swa->buffer) { + res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens; + //res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there + } + + if (self_kq_mask_swa && self_kq_mask_swa->buffer) { + res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams); + } + + return res; +} + + +void llm_graph_input_attn_k_iswa::set_input(const llama_ubatch * ubatch) { + // base tensors may not be allocated if there are no non-SWA attention layers + if (self_k_idxs && self_k_idxs->buffer) { + mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch); + } + + // the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live + if (self_kq_mask && self_kq_mask->buffer) { + mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + } + + // swa tensors may not be allocated if there are no SWA attention layers + if (self_k_idxs_swa && self_k_idxs_swa->buffer) { + mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch); + } + + if (self_kq_mask_swa && self_kq_mask_swa->buffer) { + mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); + } + + if (self_k_rot && self_k_rot->buffer) { + mctx->get_base()->set_input_k_rot(self_k_rot); + } + + if (self_k_rot_swa && self_k_rot_swa->buffer) { + mctx->get_swa()->set_input_k_rot(self_k_rot_swa); + } +} + +bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) { + const auto * mctx = static_cast(params.mctx); + + this->mctx = mctx; + + bool res = true; + + // base tensors may not be allocated if there are no non-SWA attention layers + if (self_k_idxs && self_k_idxs->buffer) { + res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; + } + + if (self_kq_mask && self_kq_mask->buffer) { + res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams); + } + + // swa tensors may not be allocated if there are no SWA attention layers + if (self_k_idxs_swa && self_k_idxs_swa->buffer) { + res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens; + } + + if (self_kq_mask_swa && self_kq_mask_swa->buffer) { + res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams); + } + + return res; +} + static void dsv4_set_i64(ggml_tensor * dst, const std::vector & src) { if (!dst || !dst->buffer) { return; } + GGML_ASSERT(dst->ne[0] == (int64_t) src.size()); - memcpy(dst->data, src.data(), src.size() * sizeof(int64_t)); + ggml_backend_tensor_set(dst, src.data(), 0, src.size()*ggml_element_size(dst)); } static void dsv4_set_i32(ggml_tensor * dst, const std::vector & src) { if (!dst || !dst->buffer) { return; } + GGML_ASSERT(dst->ne[0] == (int64_t) src.size()); - memcpy(dst->data, src.data(), src.size() * sizeof(int32_t)); + ggml_backend_tensor_set(dst, src.data(), 0, src.size()*ggml_element_size(dst)); } static void dsv4_set_kq_mask( @@ -610,18 +744,21 @@ static void dsv4_set_kq_mask( return; } - GGML_ASSERT(dst->ne[0] == plan.n_kv); - GGML_ASSERT(dst->ne[1] == (int64_t) n_tokens / n_stream); - GGML_ASSERT(dst->ne[2] == n_stream); GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); GGML_ASSERT(n_stream > 0); GGML_ASSERT(n_tokens%n_stream == 0); + GGML_ASSERT(dst->ne[0] == plan.n_kv); + GGML_ASSERT(dst->ne[1] == (int64_t) n_tokens/n_stream); + GGML_ASSERT(dst->ne[2] == 1); + GGML_ASSERT(dst->ne[3] == n_stream); + GGML_ASSERT((int64_t) plan.n_visible.size() == (int64_t) n_tokens); + GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); if (dst->type == GGML_TYPE_F32) { float * data = (float *) dst->data; for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { - const int32_t n_visible = plan.n_visible[i / (n_tokens / n_stream)]; + const int32_t n_visible = plan.n_visible[i]; for (int64_t j = 0; j < dst->ne[0]; ++j) { data[i*dst->ne[0] + j] = j < n_visible ? 0.0f : -INFINITY; @@ -633,7 +770,7 @@ static void dsv4_set_kq_mask( const ggml_fp16_t fp16_zero = llama_cast(0.0f); for (int64_t i = 0; i < (int64_t) n_tokens; ++i) { - const int32_t n_visible = plan.n_visible[i / (n_tokens / n_stream)]; + const int32_t n_visible = plan.n_visible[i]; for (int64_t j = 0; j < dst->ne[0]; ++j) { data[i*dst->ne[0] + j] = j < n_visible ? fp16_zero : fp16_ninf; @@ -642,11 +779,54 @@ static void dsv4_set_kq_mask( } } +static ggml_tensor * dsv4_build_raw_kq_mask( + ggml_context * ctx, + const llama_kv_cache_dsv4_raw_context * mctx, + const llama_ubatch & ubatch, + const llama_cparams & cparams, + int64_t n_stream) { + const auto n_kv = mctx->get_n_kv(); + const auto n_tokens = ubatch.n_tokens; + + GGML_ASSERT(n_stream > 0); + GGML_ASSERT(n_tokens%n_stream == 0); + + const auto type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; + + ggml_tensor * res = ggml_new_tensor_4d(ctx, type, n_kv, n_tokens/n_stream, 1, n_stream); + ggml_set_input(res); + ggml_set_name(res, "attn_inp_kq_mask"); + + return res; +} + +static bool dsv4_can_reuse_raw_kq_mask( + ggml_tensor * kq_mask, + const llama_kv_cache_dsv4_raw_context * mctx, + const llama_ubatch & ubatch, + int64_t n_stream) { + const auto n_kv = mctx->get_n_kv(); + const auto n_tokens = ubatch.n_tokens; + + GGML_ASSERT(n_stream > 0); + + bool res = true; + + res &= (kq_mask->ne[0] == n_kv); + res &= (kq_mask->ne[1] == n_tokens/n_stream); + res &= (kq_mask->ne[2] == 1); + res &= (kq_mask->ne[3] == n_stream); + + return res; +} + static std::string dsv4_plan_positions(const std::vector & values) { std::ostringstream ss; ss << "["; for (size_t i = 0; i < values.size(); ++i) { - if (i > 0) ss << ", "; + if (i > 0) { + ss << ", "; + } ss << values[i]; } ss << "]"; @@ -658,6 +838,7 @@ static bool dsv4_compress_debug() { const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG"); return env && atoi(env) > 0; }(); + return debug; } @@ -671,6 +852,10 @@ static void dsv4_set_comp_inputs( dsv4_set_i32(inp.state_pos, plan.state_pos); dsv4_set_i32(inp.state_persist_src_idxs, plan.state_persist_src_idxs); dsv4_set_i32(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs); + dsv4_set_i32(inp.state_restore_src_idxs, plan.state_restore_src_idxs); + dsv4_set_i32(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs); + dsv4_set_i32(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs); + dsv4_set_i32(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs); dsv4_set_i32(inp.state_read_idxs, plan.state_read_idxs); dsv4_set_i64(inp.state_write_idxs, plan.state_write_idxs); dsv4_set_i32(inp.state_write_pos, plan.state_write_pos); @@ -705,6 +890,7 @@ static bool dsv4_can_reuse_kq_mask( t->ne[2] == 1 && t->ne[3] == n_stream; } + static bool dsv4_can_reuse_comp_input( const llm_graph_input_dsv4::comp_input & inp, const llama_kv_cache_dsv4_context::comp_plan & plan, @@ -714,6 +900,10 @@ static bool dsv4_can_reuse_comp_input( res &= dsv4_can_reuse_tensor_1d(inp.state_pos, plan.state_pos.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_persist_src_idxs, plan.state_persist_src_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_persist_dst_idxs, plan.state_persist_dst_idxs.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_restore_src_idxs, plan.state_restore_src_idxs.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_restore_dst_idxs, plan.state_restore_dst_idxs.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_src_idxs, plan.state_snapshot_src_idxs.size()); + res &= dsv4_can_reuse_tensor_1d(inp.state_snapshot_dst_idxs, plan.state_snapshot_dst_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_read_idxs, plan.state_read_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_write_idxs, plan.state_write_idxs.size()); res &= dsv4_can_reuse_tensor_1d(inp.state_write_pos, plan.state_write_pos.size()); @@ -722,54 +912,6 @@ static bool dsv4_can_reuse_comp_input( return res; } -static ggml_tensor * dsv4_build_raw_kq_mask( - ggml_context * ctx, - const llama_kv_cache_dsv4_raw_context * mctx, - const llama_ubatch & ubatch, - const llama_cparams & cparams, - int64_t n_stream) { - if (n_stream == 0) { - return nullptr; - } - - const uint32_t n_tokens = ubatch.n_tokens; - const uint32_t n_kv = mctx->get_n_kv(); - - GGML_ASSERT(n_stream > 0); - GGML_ASSERT(n_tokens%n_stream == 0); - - const auto type = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; - - auto result = ggml_new_tensor_4d(ctx, type, n_kv, n_tokens/n_stream, 1, n_stream); - ggml_set_name(result, "attn_inp_kq_mask"); - ggml_set_input(result); - - return result; -} -static bool dsv4_can_reuse_raw_kq_mask( - ggml_tensor * kq_mask, - const llama_kv_cache_dsv4_raw_context * mctx, - const llama_ubatch & ubatch, - int64_t n_stream) { - const auto n_kv = mctx->get_n_kv(); - const auto n_tokens = ubatch.n_tokens; - - if (n_kv == 0) { - return true; - } - - GGML_ASSERT(n_stream > 0); - - bool res = true; - - res &= (kq_mask->ne[0] == n_kv); - res &= (kq_mask->ne[1] == n_tokens/n_stream); - res &= (kq_mask->ne[2] == 1); - res &= (kq_mask->ne[3] == n_stream); - - return res; -} - static ggml_tensor * dsv4_build_input_1d( ggml_context * ctx, ggml_type type, @@ -779,11 +921,11 @@ static ggml_tensor * dsv4_build_input_1d( return nullptr; } - auto result = ggml_new_tensor_1d(ctx, type, ne0); - ggml_set_name(result, name.c_str()); - ggml_set_input(result); + ggml_tensor * res = ggml_new_tensor_1d(ctx, type, ne0); + ggml_set_input(res); + ggml_set_name(res, name.c_str()); - return result; + return res; } static void dsv4_build_comp_inputs( @@ -796,9 +938,14 @@ static void dsv4_build_comp_inputs( inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos"); inp.state_persist_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_src_idxs.size(), std::string("dsv4_") + name + "_state_persist_src_idxs"); inp.state_persist_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_persist_dst_idxs.size(), std::string("dsv4_") + name + "_state_persist_dst_idxs"); + inp.state_restore_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_src_idxs.size(), std::string("dsv4_") + name + "_state_restore_src_idxs"); + inp.state_restore_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_restore_dst_idxs.size(), std::string("dsv4_") + name + "_state_restore_dst_idxs"); + inp.state_snapshot_src_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_src_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_src_idxs"); + inp.state_snapshot_dst_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_snapshot_dst_idxs.size(), std::string("dsv4_") + name + "_state_snapshot_dst_idxs"); inp.state_read_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_read_idxs.size(), std::string("dsv4_") + name + "_state_read_idxs"); inp.state_write_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I64, plan.state_write_idxs.size(), std::string("dsv4_") + name + "_state_write_idxs"); inp.state_write_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_write_pos.size(), std::string("dsv4_") + name + "_state_write_pos"); + if (plan.n_kv > 0) { const int64_t n_tokens = (int64_t) plan.n_visible.size(); @@ -809,7 +956,6 @@ static void dsv4_build_comp_inputs( ggml_set_input(inp.kq_mask); ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str()); } - } void llm_graph_input_dsv4_raw::set_input(const llama_ubatch * ubatch) { @@ -881,78 +1027,7 @@ bool llm_graph_input_dsv4::can_reuse(const llm_graph_params & params) { return res; } -void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) { - // base tensors may not be allocated if there are no non-SWA attention layers - if (self_k_idxs && self_k_idxs->buffer) { - mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch); - if (self_v_idxs) { - mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch); - } - } - - // the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live - if (self_kq_mask && self_kq_mask->buffer) { - mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); - } - - // swa tensors may not be allocated if there are no SWA attention layers - if (self_k_idxs_swa && self_k_idxs_swa->buffer) { - mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch); - if (self_v_idxs_swa) { - mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch); - } - } - - if (self_kq_mask_swa && self_kq_mask_swa->buffer) { - mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); - } - - if (self_k_rot && self_k_rot->buffer) { - mctx->get_base()->set_input_k_rot(self_k_rot); - } - - if (self_v_rot && self_v_rot->buffer) { - mctx->get_base()->set_input_v_rot(self_v_rot); - } - - if (self_k_rot_swa && self_k_rot_swa->buffer) { - mctx->get_swa()->set_input_k_rot(self_k_rot_swa); - } - - if (self_v_rot_swa && self_v_rot_swa->buffer) { - mctx->get_swa()->set_input_v_rot(self_v_rot_swa); - } -} - -bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { - const auto * mctx = static_cast(params.mctx); - - this->mctx = mctx; - - bool res = true; - - // base tensors may not be allocated if there are no non-SWA attention layers - if (self_k_idxs && self_k_idxs->buffer) { - res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; - //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there - } - - if (self_kq_mask && self_kq_mask->buffer) { - res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams); - } - - // swa tensors may not be allocated if there are no SWA attention layers - if (self_k_idxs_swa && self_k_idxs_swa->buffer) { - res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens; - //res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there - } - if (self_kq_mask_swa && self_kq_mask_swa->buffer) { - res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams); - } - - return res; -} void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) { GGML_ASSERT(cross_kq_mask); @@ -1231,7 +1306,7 @@ void llm_graph_result::reset() { t_embd_pooled = nullptr; t_h_nextn = nullptr; - t_layer_inp.resize(LLAMA_MAX_LAYERS); + t_layer_inp.resize(LLAMA_MAX_LAYERS + 1); std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr); t_sampled.clear(); @@ -1671,7 +1746,7 @@ ggml_tensor * llm_graph_context::build_ffn( tmp = ggml_clamp(ctx0, tmp, -limit, limit); cb(tmp, "ffn_up_clamped", il); - if (arch == LLM_ARCH_DEEPSEEK4) { + if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_clamp(ctx0, cur, -INFINITY, limit); cb(cur, "ffn_gate_clamped", il); cur = ggml_swiglu_split(ctx0, cur, tmp); @@ -2068,7 +2143,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( up = ggml_clamp(ctx0, up, -limit, limit); cb(up, "ffn_moe_up_clamped", il); - if (arch == LLM_ARCH_DEEPSEEK4) { + if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_clamp(ctx0, cur, -INFINITY, limit); cb(cur, "ffn_moe_gate_clamped", il); cur = ggml_swiglu_split(ctx0, cur, up); @@ -3135,6 +3210,75 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } +ggml_tensor * llm_graph_context::build_attn( + llm_graph_input_attn_k_iswa * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * kq_b, + ggml_tensor * sinks, + ggml_tensor * v_mla, + float kq_scale, + int il) const { + const bool is_swa = hparams.is_swa(il); + + GGML_UNUSED(v_cur); + + auto * k_rot = is_swa ? inp->self_k_rot_swa : inp->self_k_rot; + + if (k_rot) { + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, k_rot); + if (k_cur) { + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, k_rot); + } + } + + // these nodes are added to the graph together so that they are not reordered + // by doing so, the number of splits in the graph is reduced + ggml_build_forward_expand(gf, q_cur); + + if (k_cur) { + ggml_build_forward_expand(gf, k_cur); + } + + const auto * mctx_iswa = inp->mctx; + const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base(); + + // optionally store to KV cache + if (k_cur) { + const auto & k_idxs = is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + } + + const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask(); + + // MLA-style attention: the cached K is used as V + ggml_tensor * q = q_cur; + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = k; + + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + cb(cur, "kqv_out", il); + + if (k_rot) { + cur = llama_mul_mat_hadamard(ctx0, cur, k_rot); + } + + if (wo) { + cur = build_lora_mm(wo, cur, wo_s); + } + + if (wo_b) { + cur = ggml_add(ctx0, cur, wo_b); + } + + return cur; +} + llm_graph_input_attn_cross * llm_graph_context::build_attn_inp_cross() const { auto inp = std::make_unique(cross); @@ -3257,6 +3401,34 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp)); } +llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { + const auto * mctx_cur = static_cast(mctx); + + auto inp = std::make_unique(hparams, cparams, mctx_cur); + + { + inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch); + + inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams); + inp->self_kq_mask_cnv = inp->self_kq_mask; + } + + { + GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA"); + + inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); + + inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); + inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa; + } + + inp->self_k_rot = mctx_cur->get_base()->build_input_k_rot(ctx0); + + inp->self_k_rot_swa = mctx_cur->get_swa()->build_input_k_rot(ctx0); + + return (llm_graph_input_attn_k_iswa *) res->add_input(std::move(inp)); +} + llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { const auto * mctx_cur = static_cast(mctx); const auto * raw_ctx = mctx_cur->get_raw(); diff --git a/src/llama-graph.h b/src/llama-graph.h index 11a512325fb..bc3553d6eb9 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -473,6 +473,45 @@ class llm_graph_input_attn_kv_iswa : public llm_graph_input_i { const llama_kv_cache_iswa_context * mctx; }; +class llm_graph_input_attn_k_iswa : public llm_graph_input_i { +public: + llm_graph_input_attn_k_iswa( + const llama_hparams & hparams, + const llama_cparams & cparams, + const llama_kv_cache_iswa_context * mctx) : + hparams(hparams), + cparams(cparams), + mctx(mctx) { + } + ~llm_graph_input_attn_k_iswa() = default; + + void set_input(const llama_ubatch * ubatch) override; + + bool can_reuse(const llm_graph_params & params) override; + + ggml_tensor * get_k_idxs() const { return self_k_idxs; } + ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; } + + ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; } + ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; } + + ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch] + ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch] + + ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + + ggml_tensor * self_k_rot = nullptr; + ggml_tensor * self_k_rot_swa = nullptr; + + const llama_hparams hparams; + const llama_cparams cparams; + + const llama_kv_cache_iswa_context * mctx; +}; + // DSV4 raw graph inputs are SWA-only, but their mask may be stream-shaped // so raw K can be concatenated with DSV4 compressed K in one attention op. class llm_graph_input_dsv4_raw { @@ -507,6 +546,10 @@ class llm_graph_input_dsv4 : public llm_graph_input_i { ggml_tensor * state_pos = nullptr; // I32 [n_state] ggml_tensor * state_persist_src_idxs = nullptr; // I32 [n_state_persist] ggml_tensor * state_persist_dst_idxs = nullptr; // I32 [n_state_persist] + ggml_tensor * state_restore_src_idxs = nullptr; // I32 [n_state_restore] + ggml_tensor * state_restore_dst_idxs = nullptr; // I32 [n_state_restore] + ggml_tensor * state_snapshot_src_idxs = nullptr; // I32 [n_state_snapshot] + ggml_tensor * state_snapshot_dst_idxs = nullptr; // I32 [n_state_snapshot] ggml_tensor * state_read_idxs = nullptr; // I32 [ratio*n_state_write] ggml_tensor * state_write_idxs = nullptr; // I64 [n_state_write] ggml_tensor * state_write_pos = nullptr; // I32 [n_state_write] @@ -1063,7 +1106,7 @@ struct llm_graph_context { ggml_tensor * build_attn_mha( ggml_tensor * q, // [n_embd_head_q, n_head_q, n_tokens] ggml_tensor * k, // [n_embd_head_k, n_head_k, n_tokens] - ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans == false) + ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans = false) ggml_tensor * kq_b, ggml_tensor * kq_mask, ggml_tensor * sinks, // [n_head_q] @@ -1155,6 +1198,24 @@ struct llm_graph_context { float kq_scale, int il) const; + llm_graph_input_attn_k_iswa * build_attn_inp_k_iswa() const; + + // note: if k_cur is not provided, it will not be stored in the memory + // note: the K cache is used as V (MLA-style attention) + ggml_tensor * build_attn( + llm_graph_input_attn_k_iswa * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens] + ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] optional + ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] optional + ggml_tensor * kq_b, + ggml_tensor * sinks, // [n_head_q] + ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] + float kq_scale, + int il) const; + llm_graph_input_attn_cross * build_attn_inp_cross() const; ggml_tensor * build_attn( diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 069da45f4ea..5caa05e8b07 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -252,7 +252,8 @@ static void dsv4_state_write_tensor_streams( uint32_t tensor_rows, uint32_t n_rows, uint32_t s0, - uint32_t ns) { + uint32_t ns, + const std::vector * stream_ids = nullptr) { const int32_t type_i = (int32_t) tensor->type; const uint64_t ne0 = tensor->ne[0]; const uint64_t rows = n_rows; @@ -273,8 +274,16 @@ static void dsv4_state_write_tensor_streams( return; } + if (stream_ids && stream_ids->size() != ns) { + throw std::runtime_error("DSV4 state tensor stream map size mismatch"); + } + for (uint32_t s = 0; s < ns; ++s) { - const size_t offset = (size_t) (s0 + s)*stream_stride; + const uint32_t stream = stream_ids ? (*stream_ids)[s] : s0 + s; + if ((int64_t) stream >= tensor->ne[2]) { + throw std::runtime_error("DSV4 state tensor stream out of range"); + } + const size_t offset = (size_t) stream*stream_stride; io.write_tensor(tensor, offset, size); } } @@ -421,7 +430,9 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( bool overlap, uint32_t state_size, uint32_t kv_size, - uint32_t n_stream) { + uint32_t n_stream, + uint32_t n_rs_seq, + const std::vector & rs_idx) { llama_kv_cache_dsv4_context::comp_plan plan; plan.n_visible.resize(ubatch.n_tokens); plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream); @@ -451,6 +462,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( std::vector overlap_cur_reads; std::map, int64_t> curr_token_idx_map; + std::map state_write_counts; for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) { @@ -513,6 +525,7 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( plan.state_write_idxs.push_back(cache_off + pos/ratio); plan.state_write_pos.push_back((int32_t) source_start); + ++state_write_counts[seq_id]; if (overlap) { const llama_pos prev_start = source_start - ratio; @@ -531,33 +544,57 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( } } - if (ratio == DSV4_CSA_RATIO && plan.state_write_idxs.empty() && !plan.state_pos.empty()) { - // Non-boundary CSA steps still need a write op so their graph matches - // boundary steps. Use a padded scratch row that is masked from attention. + if (ratio == DSV4_CSA_RATIO && !plan.state_pos.empty()) { assert(kv_size > 0); - uint32_t i = 0; - while (i < ubatch.n_tokens && ubatch.pos[i] < 0) { - ++i; - } - assert(i < ubatch.n_tokens); + // Pad each stream to the reserve plan's block count. + const auto append_dummy_block = [&](llama_seq_id seq_id, uint32_t i) { + const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size); + const int32_t source_idx = state_source_idx(seq_id, ubatch.pos[i]); - const llama_pos pos = ubatch.pos[i]; - const llama_seq_id seq_id = ubatch.seq_id[i][0]; - const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size); - const int32_t source_idx = state_source_idx(seq_id, pos); + plan.state_write_idxs.push_back(cache_off + kv_size - 1); + plan.state_write_pos .push_back(0); - plan.state_write_idxs.push_back(cache_off + kv_size - 1); - plan.state_write_pos .push_back(0); + if (overlap) { + for (uint32_t j = 0; j < ratio; ++j) { + overlap_prev_reads.push_back(source_idx); + overlap_cur_reads .push_back(source_idx); + } + } else { + for (uint32_t j = 0; j < ratio; ++j) { + plan.state_read_idxs.push_back(source_idx); + } + } + }; - if (overlap) { - for (uint32_t j = 0; j < ratio; ++j) { - overlap_prev_reads.push_back(source_idx); - overlap_cur_reads .push_back(source_idx); + if (dsv4_ubatch_has_coupled(ubatch)) { + if (plan.state_write_idxs.empty()) { + uint32_t i = 0; + while (i < ubatch.n_tokens && ubatch.pos[i] < 0) { + ++i; + } + assert(i < ubatch.n_tokens); + append_dummy_block(ubatch.seq_id[i][0], i); } } else { - for (uint32_t j = 0; j < ratio; ++j) { - plan.state_read_idxs.push_back(source_idx); + const uint32_t n_blocks = (std::max(1, ubatch.n_seq_tokens) + ratio - 1)/ratio; + + for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { + const llama_seq_id seq_id = ubatch.seq_id_unq[s]; + const uint32_t n_writes = state_write_counts[seq_id]; + if (n_writes >= n_blocks) { + continue; + } + if (n_writes + 1 != n_blocks) { + throw std::runtime_error("DSV4 CSA sequence positions are not contiguous"); + } + + uint32_t i = 0; + while (i < ubatch.n_tokens && (ubatch.pos[i] < 0 || !dsv4_token_has_seq(ubatch, i, seq_id))) { + ++i; + } + assert(i < ubatch.n_tokens); + append_dummy_block(seq_id, i); } } } @@ -583,6 +620,63 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_comp_plan( plan.state_persist_dst_idxs.push_back(row.dst); } + + if (n_rs_seq > 0) { + for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { + const llama_seq_id seq_id = ubatch.seq_id_unq[s]; + if (seq_id < 0 || (uint32_t) seq_id >= n_stream) { + continue; + } + + const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size); + const uint32_t rollback = (uint32_t) seq_id < rs_idx.size() ? rs_idx[seq_id] : 0; + // Keep the restore graph fixed-width when no rollback is pending. + const int64_t src_plane = rollback > 0 && rollback <= n_rs_seq ? (int64_t) rollback*state_rows : 0; + for (uint32_t r = 0; r < state_size; ++r) { + plan.state_restore_src_idxs.push_back((int32_t) (src_plane + stream_off + r)); + plan.state_restore_dst_idxs.push_back((int32_t) (stream_off + r)); + } + + std::vector token_idxs; + token_idxs.reserve(ubatch.n_tokens); + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (dsv4_token_has_seq(ubatch, i, seq_id)) { + token_idxs.push_back(i); + } + } + if (token_idxs.empty()) { + continue; + } + + const uint32_t n_seq_tokens = (uint32_t) token_idxs.size(); + const int64_t scratch_off = (int64_t) state_rows*(1 + n_rs_seq); + for (uint32_t d = 1; d <= n_rs_seq; ++d) { + const int64_t dst_plane = (int64_t) d*state_rows; + + for (uint32_t r = 0; r < state_size; ++r) { + int32_t src; + if (d <= n_seq_tokens) { + const uint32_t prefix = n_seq_tokens - d; + src = (int32_t) (stream_off + r); + + for (uint32_t j = 0; j < prefix; ++j) { + const uint32_t i_tok = token_idxs[j]; + if (ubatch.pos[i_tok] >= 0 && (uint32_t) (ubatch.pos[i_tok]%state_size) == r) { + src = (int32_t) (scratch_off + i_tok); + } + } + } else { + const int64_t src_plane = (int64_t) (d - n_seq_tokens)*state_rows; + src = (int32_t) (src_plane + stream_off + r); + } + + plan.state_snapshot_src_idxs.push_back(src); + plan.state_snapshot_dst_idxs.push_back((int32_t) (dst_plane + stream_off + r)); + } + } + } + } + static const bool debug = []() { const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG"); return env && atoi(env) > 0; @@ -604,12 +698,14 @@ static std::vector dsv4_build_comp_plans bool overlap, uint32_t state_size, uint32_t kv_size, - uint32_t n_stream) { + uint32_t n_stream, + uint32_t n_rs_seq, + const std::vector & rs_idx) { std::vector plans; plans.reserve(ubatches.size()); for (const llama_ubatch & ubatch : ubatches) { - plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream)); + plans.push_back(dsv4_build_comp_plan(ubatch, ratio, overlap, state_size, kv_size, n_stream, n_rs_seq, rs_idx)); } return plans; @@ -696,7 +792,8 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan( bool overlap, uint32_t state_size, uint32_t kv_size, - uint32_t n_stream) { + uint32_t n_stream, + uint32_t n_rs_seq) { llama_kv_cache_dsv4_context::comp_plan plan; plan.n_visible.resize(ubatch.n_tokens); plan.n_stream = dsv4_comp_graph_n_stream(ubatch, n_stream); @@ -714,10 +811,16 @@ static llama_kv_cache_dsv4_context::comp_plan dsv4_build_reserve_comp_plan( const uint64_t state_rows = (uint64_t) state_size*n_stream; const size_t n_persist = (size_t) std::min(ubatch.n_tokens, state_rows); + const size_t n_restore = n_rs_seq > 0 ? (size_t) state_size*std::max(1, ubatch.n_seqs_unq) : 0; + const size_t n_snapshot = (size_t) n_rs_seq*state_size*std::max(1, ubatch.n_seqs_unq); plan.state_pos .resize(ubatch.n_tokens); plan.state_persist_src_idxs.resize(n_persist); plan.state_persist_dst_idxs.resize(n_persist); + plan.state_restore_src_idxs.resize(n_restore); + plan.state_restore_dst_idxs.resize(n_restore); + plan.state_snapshot_src_idxs.resize(n_snapshot); + plan.state_snapshot_dst_idxs.resize(n_snapshot); plan.state_read_idxs .resize((overlap ? 2u : 1u)*ratio*n_blocks); plan.state_write_idxs.resize(n_blocks); plan.state_write_pos .resize(n_blocks); @@ -743,12 +846,14 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( uint32_t ratio, uint32_t state_size, uint32_t n_embd_state, + uint32_t n_rs_seq, const char * name, const llama_memory_i::layer_filter_cb & filter) : ratio(ratio), state_size(state_size), n_embd_state(n_embd_state), - n_stream(unified ? 1 : n_seq_max) { + n_stream(unified ? 1 : n_seq_max), + n_rs_seq(n_rs_seq) { const llama_hparams & hparams = model.hparams; struct ggml_backend_buft_comparator { @@ -804,8 +909,9 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( throw std::runtime_error("failed to create ggml context for DSV4 compressor state"); } - ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream); - ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_stream); + const uint32_t n_planes = n_stream*(1 + n_rs_seq); + ggml_tensor * kv = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes); + ggml_tensor * score = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_state, state_size, n_planes); ggml_format_name(kv, "dsv4_%s_state_kv_l%d", name, il); ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il); @@ -837,8 +943,8 @@ llama_dsv4_comp_state::llama_dsv4_comp_state( ctxs_bufs.emplace_back(std::move(ctx), buf); } - LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, layers = %zu, size = %7.2f MiB\n", - __func__, name, ratio, state_size, n_embd_state, n_stream, layers.size(), total_size()/1024.0/1024.0); + LLAMA_LOG_INFO("%s: %s ratio = %u, state = %u x %u, streams = %u, rs_seq = %u, layers = %zu, size = %7.2f MiB\n", + __func__, name, ratio, state_size, n_embd_state, n_stream, n_rs_seq, layers.size(), total_size()/1024.0/1024.0); } void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) { @@ -848,9 +954,13 @@ void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) { if (seq_id >= 0) { GGML_ASSERT((uint32_t) seq_id < n_stream); + for (const auto & layer : layers) { - dsv4_clear_tensor_stream(layer.kv, (uint32_t) seq_id); - dsv4_clear_tensor_stream(layer.score, (uint32_t) seq_id); + for (uint32_t d = 0; d <= n_rs_seq; ++d) { + const uint32_t stream = d*n_stream + (uint32_t) seq_id; + dsv4_clear_tensor_stream(layer.kv, stream); + dsv4_clear_tensor_stream(layer.score, stream); + } } return; } @@ -868,6 +978,8 @@ void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ return; } + clear(seq_id_dst, true); + sc_info.ssrc.push_back((uint32_t) seq_id_src); sc_info.sdst.push_back((uint32_t) seq_id_dst); } @@ -896,6 +1008,14 @@ uint32_t llama_dsv4_comp_state::get_n_stream() const { return n_stream; } +uint32_t llama_dsv4_comp_state::get_n_rs_seq() const { + return n_rs_seq; +} + +uint32_t llama_dsv4_comp_state::get_n_rows() const { + return state_size*n_stream; +} + std::map llama_dsv4_comp_state::memory_breakdown() const { std::map ret; for (const auto & [_, buf] : ctxs_bufs) { @@ -905,13 +1025,26 @@ std::map llama_dsv4_comp_state::memory_break return ret; } -void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { +void llama_dsv4_comp_state::state_write( + llama_io_write_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags, + const std::vector & rs_idx) const { GGML_UNUSED(flags); uint32_t s0; uint32_t ns; dsv4_state_src_stream_range(n_stream, seq_id, s0, ns); + std::vector stream_ids(ns); + for (uint32_t s = 0; s < ns; ++s) { + const uint32_t seq = seq_id >= 0 ? (uint32_t) seq_id : s0 + s; + if (seq >= rs_idx.size() || rs_idx[seq] > n_rs_seq) { + throw std::runtime_error("DSV4 recurrent state rollback index out of range"); + } + stream_ids[s] = rs_idx[seq]*n_stream + s0 + s; + } + const uint32_t version = DSV4_COMP_STATE_VER; const uint32_t n_layer = layers.size(); @@ -925,8 +1058,8 @@ void llama_dsv4_comp_state::state_write(llama_io_write_i & io, llama_seq_id seq_ for (const auto & layer : layers) { io.write(&layer.il, sizeof(layer.il)); - dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns); - dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns); + dsv4_state_write_tensor_streams(io, layer.kv, state_size, state_size, s0, ns, &stream_ids); + dsv4_state_write_tensor_streams(io, layer.score, state_size, state_size, s0, ns, &stream_ids); } } @@ -972,28 +1105,40 @@ void llama_dsv4_comp_state::state_read(llama_io_read_i & io, llama_seq_id seq_id } } -ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const { +ggml_tensor * llama_dsv4_comp_state::get_kv_all(ggml_context * ctx, int32_t il) const { const int32_t ids = map_layer_ids.at(il); - ggml_tensor * state = layers[ids].kv; - return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]); + return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0); } -ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const { +ggml_tensor * llama_dsv4_comp_state::get_score_all(ggml_context * ctx, int32_t il) const { const int32_t ids = map_layer_ids.at(il); - ggml_tensor * state = layers[ids].score; - return ggml_reshape_2d(ctx, state, state->ne[0], state->ne[1]*state->ne[2]); + return ggml_view_2d(ctx, state, state->ne[0], get_n_rows()*(1 + n_rs_seq), state->nb[1], 0); +} + +ggml_tensor * llama_dsv4_comp_state::get_kv(ggml_context * ctx, int32_t il) const { + ggml_tensor * state = get_kv_all(ctx, il); + const size_t row_size = ggml_row_size(state->type, state->ne[0]); + + return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size); +} + +ggml_tensor * llama_dsv4_comp_state::get_score(ggml_context * ctx, int32_t il) const { + ggml_tensor * state = get_score_all(ctx, il); + const size_t row_size = ggml_row_size(state->type, state->ne[0]); + + return ggml_view_2d(ctx, state, state->ne[0], get_n_rows(), state->nb[1], 0*row_size); } ggml_tensor * llama_dsv4_comp_state::cpy_kv(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const { - return ggml_set_rows(ctx, get_kv(ctx, il), cur, idxs); + return ggml_set_rows(ctx, get_kv_all(ctx, il), cur, idxs); } ggml_tensor * llama_dsv4_comp_state::cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const { - return ggml_set_rows(ctx, get_score(ctx, il), cur, idxs); + return ggml_set_rows(ctx, get_score_all(ctx, il), cur, idxs); } size_t llama_dsv4_comp_state::total_size() const { @@ -1022,13 +1167,16 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( uint32_t n_seq_max, uint32_t n_ubatch, uint32_t n_pad, + uint32_t n_rs_seq, const layer_filter_cb & filter, const layer_reuse_cb & reuse) : hparams_raw(model.hparams), hparams_csa(model.hparams), hparams_hca(model.hparams), hparams_lid(model.hparams), - n_seq_max(n_seq_max) { + n_seq_max(n_seq_max), + n_rs_seq(n_rs_seq), + rs_idx(n_seq_max, 0) { const layer_filter_cb filter_raw = [&](int32_t il) { if (filter && !filter(il)) { @@ -1043,6 +1191,11 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( // Keep DSV4 KV/state streams per sequence even when public KV mode is unified. const bool unified_raw = false; + hparams_raw.n_layer_nextn = 0; + hparams_csa.n_layer_nextn = 0; + hparams_hca.n_layer_nextn = 0; + hparams_lid.n_layer_nextn = 0; + LLAMA_LOG_INFO("%s: creating DSV4 raw KV cache\n", __func__); dsv4_make_k_only(hparams_raw); @@ -1109,19 +1262,19 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( csa_state = std::make_unique( model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, - 2*model.hparams.n_embd_head_k(), "csa", filter_csa); + 2*model.hparams.n_embd_head_k(), n_rs_seq, "csa", filter_csa); LLAMA_LOG_INFO("%s: creating DSV4 HCA compressor state\n", __func__); hca_state = std::make_unique( model, offload, unified_compressed, n_seq_max, DSV4_HCA_RATIO, DSV4_HCA_RATIO, - model.hparams.n_embd_head_k(), "hca", filter_hca); + model.hparams.n_embd_head_k(), n_rs_seq, "hca", filter_hca); LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer compressor state\n", __func__); lid_state = std::make_unique( model, offload, unified_compressed, n_seq_max, DSV4_CSA_RATIO, 2*DSV4_CSA_RATIO, - 2*model.hparams.indexer_head_size, "lid", filter_csa); + 2*model.hparams.indexer_head_size, n_rs_seq, "lid", filter_csa); // DSV4 attention reads compressed-K / compressor-state rows that the current // graph does not necessarily overwrite; uninitialized buffer contents would @@ -1255,17 +1408,35 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 } if (p0 > 0) { - if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max || - p0 <= kv_raw->seq_pos_max(seq_id)) { + if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { + return false; + } + + const llama_pos pos_max = kv_raw->seq_pos_max(seq_id); + if (p0 > pos_max) { + bool res = true; + + res = res & kv_raw->seq_rm(seq_id, p0, -1); + res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1); + res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1); + res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1); + + return res; + } + + if (n_rs_seq == 0) { return false; } - bool res = true; + const llama_pos rollback = pos_max - (p0 - 1); + if (rollback < 1 || rollback > (llama_pos) n_rs_seq) { + return false; + } - res = res & kv_raw->seq_rm(seq_id, p0, -1); - res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1); - res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1); - res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1); + const bool res = kv_raw->seq_rm(seq_id, p0, p1); + if (res) { + rs_idx[seq_id] = (uint32_t) rollback; + } return res; } @@ -1290,6 +1461,10 @@ void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_ds csa_state->seq_cp(seq_id_src, seq_id_dst); hca_state->seq_cp(seq_id_src, seq_id_dst); lid_state->seq_cp(seq_id_src, seq_id_dst); + + if (seq_id_src != seq_id_dst) { + rs_idx[seq_id_dst] = 0; + } } void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) { @@ -1386,9 +1561,9 @@ void llama_kv_cache_dsv4::state_write(llama_io_write_i & io, llama_seq_id seq_id dsv4_state_write_k_cache(io, kv_lid.get(), seq_id, flags, n_rows_lid); } - csa_state->state_write(io, seq_id, flags); - hca_state->state_write(io, seq_id, flags); - lid_state->state_write(io, seq_id, flags); + csa_state->state_write(io, seq_id, flags, rs_idx); + hca_state->state_write(io, seq_id, flags, rs_idx); + lid_state->state_write(io, seq_id, flags, rs_idx); } void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { @@ -1432,6 +1607,12 @@ void llama_kv_cache_dsv4::state_read(llama_io_read_i & io, llama_seq_id seq_id, hca_state->state_read(io, seq_id, flags); lid_state->state_read(io, seq_id, flags); + if (seq_id >= 0) { + GGML_ASSERT((uint32_t) seq_id < n_seq_max); + rs_idx[seq_id] = 0; + } else { + std::fill(rs_idx.begin(), rs_idx.end(), 0); + } } llama_kv_cache_iswa * llama_kv_cache_dsv4::get_raw() const { @@ -1462,6 +1643,31 @@ llama_dsv4_comp_state * llama_kv_cache_dsv4::get_lid_state() const { return lid_state.get(); } +uint32_t llama_kv_cache_dsv4::get_n_rs_seq() const { + return n_rs_seq; +} + +const std::vector & llama_kv_cache_dsv4::get_rs_idx() const { + return rs_idx; +} + +void llama_kv_cache_dsv4::reset_rs_idx_for_ubatches(const std::vector & ubatches) { + if (n_rs_seq == 0) { + return; + } + + for (const llama_ubatch & ubatch : ubatches) { + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + for (int32_t s = 0; s < ubatch.n_seq_id[i]; ++s) { + const llama_seq_id seq_id = ubatch.seq_id[i][s]; + if (seq_id >= 0 && (uint32_t) seq_id < n_seq_max) { + rs_idx[seq_id] = 0; + } + } + } + } +} + void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) { if (seq_id < 0) { kv_csa->clear(data); @@ -1488,6 +1694,12 @@ void llama_kv_cache_dsv4::clear_compressed(llama_seq_id seq_id, bool data) { csa_state->clear(seq_id, data); hca_state->clear(seq_id, data); lid_state->clear(seq_id, data); + + if (seq_id >= 0) { + rs_idx[seq_id] = 0; + } else { + std::fill(rs_idx.begin(), rs_idx.end(), 0); + } } // @@ -1779,10 +1991,14 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( std::vector ubatches_raw) : ubatches(std::move(ubatches)), plans_csa(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true, - kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream())), + kv->get_csa_state()->get_state_size(), kv->get_csa()->get_size(), kv->get_csa_state()->get_n_stream(), + kv->get_n_rs_seq(), kv->get_rs_idx())), plans_hca(dsv4_build_comp_plans(this->ubatches, DSV4_HCA_RATIO, false, - kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream())), - plans_lid(plans_csa), + kv->get_hca_state()->get_state_size(), kv->get_hca()->get_size(), kv->get_hca_state()->get_n_stream(), + kv->get_n_rs_seq(), kv->get_rs_idx())), + plans_lid(dsv4_build_comp_plans(this->ubatches, DSV4_CSA_RATIO, true, + kv->get_lid_state()->get_state_size(), kv->get_lid()->get_size(), kv->get_lid_state()->get_n_stream(), + kv->get_n_rs_seq(), kv->get_rs_idx())), ctx_raw(std::make_unique( kv->get_raw(), std::move(sinfos_raw_base_write), @@ -1809,6 +2025,7 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context( hca_state(kv->get_hca_state()), lid_state(kv->get_lid_state()), status(ctx_raw->get_status()) { + kv->reset_rs_idx_for_ubatches(this->ubatches); } llama_kv_cache_dsv4_context::~llama_kv_cache_dsv4_context() = default; @@ -1944,7 +2161,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ reserve_plan_csa = dsv4_build_reserve_comp_plan( ubatch, DSV4_CSA_RATIO, true, - csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream()); + csa_state->get_state_size(), get_csa()->get_n_kv(), csa_state->get_n_stream(), csa_state->get_n_rs_seq()); return reserve_plan_csa; } @@ -1958,7 +2175,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ reserve_plan_hca = dsv4_build_reserve_comp_plan( ubatch, DSV4_HCA_RATIO, false, - hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream()); + hca_state->get_state_size(), get_hca()->get_n_kv(), hca_state->get_n_stream(), hca_state->get_n_rs_seq()); return reserve_plan_hca; } @@ -1972,7 +2189,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ reserve_plan_lid = dsv4_build_reserve_comp_plan( ubatch, DSV4_CSA_RATIO, true, - lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream()); + lid_state->get_state_size(), get_lid()->get_n_kv(), lid_state->get_n_stream(), lid_state->get_n_rs_seq()); return reserve_plan_lid; } diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index 76b1daf5787..ce39867c034 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -22,6 +22,7 @@ class llama_dsv4_comp_state { uint32_t ratio, uint32_t state_size, uint32_t n_embd_state, + uint32_t n_rs_seq, const char * name, const llama_memory_i::layer_filter_cb & filter); @@ -29,17 +30,21 @@ class llama_dsv4_comp_state { void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst); void apply_copies(const stream_copy_info & sc_info) const; - uint32_t get_ratio() const; + uint32_t get_ratio() const; uint32_t get_state_size() const; - uint32_t get_n_stream() const; + uint32_t get_n_stream() const; + uint32_t get_n_rs_seq() const; + uint32_t get_n_rows() const; std::map memory_breakdown() const; - void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const; + void state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags, const std::vector & rs_idx) const; void state_read (llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags); - ggml_tensor * get_kv (ggml_context * ctx, int32_t il) const; - ggml_tensor * get_score(ggml_context * ctx, int32_t il) const; + ggml_tensor * get_kv (ggml_context * ctx, int32_t il) const; + ggml_tensor * get_score (ggml_context * ctx, int32_t il) const; + ggml_tensor * get_kv_all (ggml_context * ctx, int32_t il) const; + ggml_tensor * get_score_all(ggml_context * ctx, int32_t il) const; ggml_tensor * cpy_kv (ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const; ggml_tensor * cpy_score(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * idxs, int32_t il) const; @@ -59,6 +64,7 @@ class llama_dsv4_comp_state { const uint32_t state_size; const uint32_t n_embd_state; const uint32_t n_stream; + const uint32_t n_rs_seq; std::vector> ctxs_bufs; @@ -93,6 +99,7 @@ class llama_kv_cache_dsv4 : public llama_memory_i { uint32_t n_seq_max, uint32_t n_ubatch, uint32_t n_pad, + uint32_t n_rs_seq, const layer_filter_cb & filter, const layer_reuse_cb & reuse); @@ -141,6 +148,10 @@ class llama_kv_cache_dsv4 : public llama_memory_i { llama_dsv4_comp_state * get_hca_state() const; llama_dsv4_comp_state * get_lid_state() const; + uint32_t get_n_rs_seq() const; + const std::vector & get_rs_idx() const; + void reset_rs_idx_for_ubatches(const std::vector & ubatches); + private: llama_hparams hparams_raw; llama_hparams hparams_csa; @@ -148,6 +159,9 @@ class llama_kv_cache_dsv4 : public llama_memory_i { llama_hparams hparams_lid; const uint32_t n_seq_max; + const uint32_t n_rs_seq; + + std::vector rs_idx; std::unique_ptr kv_raw; std::unique_ptr kv_csa; @@ -268,6 +282,17 @@ class llama_kv_cache_dsv4_context : public llama_memory_context_i { std::vector state_persist_src_idxs; std::vector state_persist_dst_idxs; + // Device-side rollback restore copies snapshot planes back to the + // current compressor-state plane before the graph reads it. + std::vector state_restore_src_idxs; + std::vector state_restore_dst_idxs; + + // Device-side rollback snapshots copy rows from the graph-local + // [persistent_state | current_ubatch_scratch] tensor into rollback + // planes after the graph has computed current-token compressor state. + std::vector state_snapshot_src_idxs; + std::vector state_snapshot_dst_idxs; + // Flattened source row ids used for state-backed commits. Source rows // index the graph-local [persistent_state | current_ubatch_scratch] // tensor. For overlapped compression the first half is previous rows diff --git a/src/llama-model.cpp b/src/llama-model.cpp index a545a7e5258..fa34791eb51 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2163,21 +2163,75 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } break; case LLM_ARCH_DEEPSEEK4: { - res = new llama_kv_cache_dsv4( - *this, - params.type_k, - params.type_v, - !cparams.flash_attn, - cparams.offload_kqv, - params.swa_full, - cparams.kv_unified, - cparams.n_ctx_seq, - cparams.n_seq_max, - cparams.n_ubatch, - 1, - nullptr, - nullptr); + + GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); + + if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) { + const llama_memory_i::layer_filter_cb filter_mtp = [&](int32_t il) { + return il >= (int32_t) hparams.n_layer(); + }; + + res = new llama_kv_cache_iswa( + *this, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + params.swa_full, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + cparams.n_ubatch, + 1, + nullptr, + filter_mtp, + nullptr, + nullptr); + } else { + res = new llama_kv_cache_dsv4( + *this, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + params.swa_full, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + cparams.n_ubatch, + 1, + cparams.n_rs_seq, + nullptr, + nullptr); + } } break; + case LLM_ARCH_DFLASH: + { + // DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring) + if (hparams.dsv4_hc_mult > 0) { + GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); + + res = new llama_kv_cache_iswa( + *this, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + params.swa_full, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + cparams.n_ubatch, + 1, + nullptr, + nullptr, + nullptr, + nullptr); + break; + } + } + [[fallthrough]]; + // Models that need standard caching should rely on recurrent/hybrid // checks default: @@ -2293,24 +2347,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } } - if (arch == LLM_ARCH_DEEPSEEK4) { - GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); - - res = new llama_kv_cache_dsv4( - *this, - params.type_k, - params.type_v, - !cparams.flash_attn, - cparams.offload_kqv, - params.swa_full, - cparams.kv_unified, - cparams.n_ctx_seq, - cparams.n_seq_max, - cparams.n_ubatch, - 1, - filter, - reuse); - } else if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { GGML_ASSERT(hparams.is_swa_any()); if (arch == LLM_ARCH_GEMMA4_ASSISTANT) { @@ -2659,9 +2696,12 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_STEP35: case LLM_ARCH_TALKIE: case LLM_ARCH_MELLUM: - case LLM_ARCH_DFLASH: return LLAMA_ROPE_TYPE_NEOX; + case LLM_ARCH_DFLASH: + // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX + return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX; + case LLM_ARCH_QWEN2VL: case LLM_ARCH_PADDLEOCR: return LLAMA_ROPE_TYPE_MROPE; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 2d41dace0b2..e68dc49b6df 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -16,6 +16,16 @@ static float dsv4_rope_attn_factor(float freq_scale, float ext_factor) { } void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + if (hparams.n_layer_nextn > 0 && hparams.n_layer_nextn < hparams.n_layer_all) { + const uint32_t n_layer_main = hparams.n_layer_all - hparams.n_layer_nextn; + const std::string mtp_probe = "blk." + std::to_string(n_layer_main) + ".nextn.eh_proj.weight"; + if (ml.get_weight(mtp_probe.c_str()) == nullptr) { + hparams.n_layer_nextn = 0; + } + } + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < block_count"); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); @@ -24,8 +34,8 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); - ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer()); - if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer(), 0)) { + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all); + if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) { hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp; } @@ -41,9 +51,11 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); ml.get_key(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count); + hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd; + uint32_t n_compress_ratios = 0; ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_compress_ratios); - if (n_compress_ratios < hparams.n_layer()) { + if (n_compress_ratios < hparams.n_layer_all) { throw std::runtime_error("DeepSeek-V4 compress_ratios is shorter than block_count"); } ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios); @@ -54,6 +66,9 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { } hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; hparams.set_swa_pattern(0); + for (uint32_t il = hparams.n_layer(); il < hparams.n_layer_all; ++il) { + hparams.is_swa_impl[il] = true; + } switch (hparams.n_layer()) { case 43: type = LLM_TYPE_UNKNOWN; break; @@ -61,7 +76,7 @@ void llama_model_deepseek4::load_arch_hparams(llama_model_loader & ml) { } } -void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) { +void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; const int64_t q_lora_rank = hparams.n_lora_q; @@ -75,6 +90,10 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) { const int64_t hc_dim = hc_mult * n_embd; const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult; + const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + const int mtp_flags = ml.load_mtp ? 0 : TENSOR_SKIP; + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); @@ -84,69 +103,82 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader &) { hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0); hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0); - for (int i = 0; i < n_layer; ++i) { + for (int i = 0; i < n_layer_all; ++i) { auto & layer = layers[i]; - - layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); - layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0); - layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0); - layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0); - layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0); - layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0); - layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0); - layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, 0); - layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0); - - layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0); - layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0); - layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0); - layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0); - layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0); - layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0); + const int flags = i < n_layer ? trunk_flags : mtp_flags; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, flags); + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, flags); + layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, flags); + layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, flags); + layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, flags); + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, flags); + + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, flags); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, flags); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, flags); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, flags); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, flags); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, flags); const int64_t ratio = hparams.dsv4_compress_ratios[i]; if (ratio != 0) { const int64_t coff = ratio == 4 ? 2 : 1; - layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, coff * n_embd_head}, 0); - layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, 0); - layer.attn_comp_ape = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE, "weight", i), {coff * n_embd_head, ratio}, 0); - layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, 0); + layer.attn_comp_wkv = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WKV, "weight", i), {n_embd, coff * n_embd_head}, flags); + layer.attn_comp_wgate = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "weight", i), {n_embd, coff * n_embd_head}, flags); + layer.attn_comp_ape = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_APE, "weight", i), {coff * n_embd_head, ratio}, flags); + layer.attn_comp_norm = create_tensor(tn(LLM_TENSOR_ATTN_COMPRESSOR_NORM, "weight", i), {n_embd_head}, flags); if (ratio == 4) { const int64_t n_embd_indexer = hparams.indexer_head_size; - layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, 0); - layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, 0); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, flags); - layer.indexer_comp_wkv = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "weight", i), {n_embd, 2 * n_embd_indexer}, 0); - layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, 0); - layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", i), {2 * n_embd_indexer, ratio}, 0); - layer.indexer_comp_norm = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "weight", i), {n_embd_indexer}, 0); + layer.indexer_comp_wkv = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "weight", i), {n_embd, 2 * n_embd_indexer}, flags); + layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, 2 * n_embd_indexer}, flags); + layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", i), {2 * n_embd_indexer, ratio}, flags); + layer.indexer_comp_norm = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "weight", i), {n_embd_indexer}, flags); } else if (ratio != 128) { throw std::runtime_error("DeepSeek-V4 loader only supports compression ratios 0, 4, and 128"); } } - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); if ((uint32_t) i < hparams.dsv4_hash_layer_count) { - layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, 0); + layer.ffn_gate_tid2eid = create_tensor(tn(LLM_TENSOR_FFN_GATE_TID2EID, "weight", i), {n_expert_used, n_vocab}, flags); } else { - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags); + } + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + + if (i >= n_layer) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags); + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED | flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED | flags); } - layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); - - layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); - layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); - - layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); - layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0); - layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); } } std::unique_ptr llama_model_deepseek4::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -175,18 +207,69 @@ static ggml_tensor * dsv4_append_zero_row(ggml_context * ctx, ggml_tensor * t, b return ggml_concat(ctx, t, row, 1); } -static ggml_tensor * dsv4_with_zero_dep(ggml_context * ctx, ggml_tensor * t, ggml_tensor * dep) { - if (dep == nullptr) { - return t; +struct dsv4_state_tensors { + ggml_tensor * kv; + ggml_tensor * score; +}; + +static dsv4_state_tensors dsv4_build_state_restore( + ggml_context * ctx, + const llm_graph_input_dsv4::comp_input & inp, + const llama_dsv4_comp_state * state, + int32_t il) { + dsv4_state_tensors restored = { + state->get_kv_all(ctx, il), + state->get_score_all(ctx, il), + }; + + if (inp.state_restore_src_idxs == nullptr || inp.state_restore_dst_idxs == nullptr) { + return restored; + } + + ggml_tensor * kv_rows = ggml_get_rows(ctx, restored.kv, inp.state_restore_src_idxs); + restored.kv = state->cpy_kv(ctx, kv_rows, inp.state_restore_dst_idxs, il); + + ggml_tensor * score_rows = ggml_get_rows(ctx, restored.score, inp.state_restore_src_idxs); + restored.score = state->cpy_score(ctx, score_rows, inp.state_restore_dst_idxs, il); + + return restored; +} + +static dsv4_state_tensors dsv4_build_state_snapshot( + ggml_context * ctx, + const llm_graph_input_dsv4::comp_input & inp, + const llama_dsv4_comp_state * state, + ggml_tensor * source_kv, + ggml_tensor * source_score, + int32_t il) { + if (inp.state_snapshot_src_idxs == nullptr || inp.state_snapshot_dst_idxs == nullptr || + source_kv == nullptr || source_score == nullptr) { + return {}; } - ggml_tensor * zero = ggml_scale(ctx, ggml_sum(ctx, dep), 0.0f); - return ggml_add(ctx, t, zero); + ggml_tensor * kv_rows = ggml_get_rows(ctx, source_kv, inp.state_snapshot_src_idxs); + ggml_tensor * kv = state->cpy_kv(ctx, kv_rows, inp.state_snapshot_dst_idxs, il); + + ggml_tensor * score_rows = ggml_get_rows(ctx, source_score, inp.state_snapshot_src_idxs); + ggml_tensor * score = state->cpy_score(ctx, score_rows, inp.state_snapshot_dst_idxs, il); + + return { kv, score }; } static constexpr int64_t DSV4_CSA_RATIO = 4; static constexpr int64_t DSV4_HCA_RATIO = 128; +// mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens] +static ggml_tensor * dsv4_hc_mean(ggml_context * ctx, ggml_tensor * x) { + const int64_t hc = x->ne[1]; + + ggml_tensor * acc = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0); + for (int64_t s = 1; s < hc; ++s) { + acc = ggml_add(ctx, acc, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1])); + } + return ggml_scale(ctx, acc, 1.0f/hc); +} + static ggml_tensor * dsv4_hc_affine( ggml_context * ctx, ggml_tensor * x, @@ -804,8 +887,29 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( ggml_tensor * cur, ggml_tensor * inp_pos, int il) const { + return build_attention_impl(model, inp_dsv4, nullptr, cur, inp_pos, il); +} + +ggml_tensor * llama_model_deepseek4::graph::build_attention( + const llama_model & model, + llm_graph_input_attn_k_iswa * inp_mtp, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int il) const { + return build_attention_impl(model, nullptr, inp_mtp, cur, inp_pos, il); +} + +ggml_tensor * llama_model_deepseek4::graph::build_attention_impl( + const llama_model & model, + llm_graph_input_dsv4 * inp_dsv4, + llm_graph_input_attn_k_iswa * inp_mtp, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int il) const { + GGML_ASSERT((inp_dsv4 == nullptr) != (inp_mtp == nullptr)); + const auto & layer = model.layers[il]; - llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw(); + llm_graph_input_dsv4_raw * inp_attn = inp_dsv4 ? inp_dsv4->get_raw() : nullptr; const int64_t n_embd_head = hparams.n_embd_head_k(); const int64_t n_embd_head_rope = hparams.n_rot(); @@ -873,9 +977,12 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( cb(kv, "kv", il); const int64_t ratio = hparams.dsv4_compress_ratios[il]; + GGML_ASSERT(inp_dsv4 || ratio == 0); ggml_tensor * hca_state_kv = nullptr; ggml_tensor * hca_state_score = nullptr; + ggml_tensor * hca_source_kv = nullptr; + ggml_tensor * hca_source_score = nullptr; if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) { hca_state_kv = build_lora_mm(layer.attn_comp_wkv, cur); cb(hca_state_kv, "hca_state_kv", il); @@ -906,10 +1013,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( GGML_ASSERT(inp_dsv4->get_csa().state_write_idxs); - ggml_tensor * csa_source_kv = ggml_concat(ctx0, - inp_dsv4->mctx->get_csa_state()->get_kv(ctx0, il), csa_state_kv, 1); - ggml_tensor * csa_source_score = ggml_concat(ctx0, - inp_dsv4->mctx->get_csa_state()->get_score(ctx0, il), csa_state_score, 1); + const auto * csa_state = inp_dsv4->mctx->get_csa_state(); + const dsv4_state_tensors csa_restored = dsv4_build_state_restore( + ctx0, inp_dsv4->get_csa(), csa_state, il); + ggml_tensor * csa_base_kv = dsv4_view_2d( + ctx0, csa_restored.kv, csa_restored.kv->ne[0], csa_state->get_n_rows(), 0); + ggml_tensor * csa_base_score = dsv4_view_2d( + ctx0, csa_restored.score, csa_restored.score->ne[0], csa_state->get_n_rows(), 0); + + ggml_tensor * csa_source_kv = ggml_concat(ctx0, csa_base_kv, csa_state_kv, 1); + ggml_tensor * csa_source_score = ggml_concat(ctx0, csa_base_score, csa_state_score, 1); ggml_tensor * kv_comp_csa_state = build_overlap_compressed_kv_from_state( csa_source_kv, @@ -930,8 +1043,19 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( ggml_build_forward_expand(gf, inp_dsv4->mctx->get_csa()->cpy_k(ctx0, kv_comp_csa_state, inp_dsv4->get_csa().state_write_idxs, il)); - csa_state_kv = dsv4_with_zero_dep(ctx0, csa_state_kv, kv_comp_csa_state); - csa_state_score = dsv4_with_zero_dep(ctx0, csa_state_score, kv_comp_csa_state); + ggml_tensor * csa_snapshot_source_kv = ggml_concat(ctx0, + csa_restored.kv, csa_state_kv, 1); + ggml_tensor * csa_snapshot_source_score = ggml_concat(ctx0, + csa_restored.score, csa_state_score, 1); + + const dsv4_state_tensors csa_snapshot = dsv4_build_state_snapshot( + ctx0, inp_dsv4->get_csa(), csa_state, csa_snapshot_source_kv, csa_snapshot_source_score, il); + if (csa_snapshot.kv != nullptr) { + ggml_build_forward_expand(gf, csa_snapshot.kv); + } + if (csa_snapshot.score != nullptr) { + ggml_build_forward_expand(gf, csa_snapshot.score); + } ggml_tensor * csa_persist_kv = ggml_get_rows(ctx0, csa_state_kv, inp_dsv4->get_csa().state_persist_src_idxs); ggml_tensor * csa_persist_score = ggml_get_rows(ctx0, csa_state_score, inp_dsv4->get_csa().state_persist_src_idxs); @@ -958,10 +1082,16 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( GGML_ASSERT(inp_dsv4->get_lid().state_write_idxs); - ggml_tensor * lid_source_kv = ggml_concat(ctx0, - inp_dsv4->mctx->get_lid_state()->get_kv(ctx0, il), lid_state_kv, 1); - ggml_tensor * lid_source_score = ggml_concat(ctx0, - inp_dsv4->mctx->get_lid_state()->get_score(ctx0, il), lid_state_score, 1); + const auto * lid_state = inp_dsv4->mctx->get_lid_state(); + const dsv4_state_tensors lid_restored = dsv4_build_state_restore( + ctx0, inp_dsv4->get_lid(), lid_state, il); + ggml_tensor * lid_base_kv = dsv4_view_2d( + ctx0, lid_restored.kv, lid_restored.kv->ne[0], lid_state->get_n_rows(), 0); + ggml_tensor * lid_base_score = dsv4_view_2d( + ctx0, lid_restored.score, lid_restored.score->ne[0], lid_state->get_n_rows(), 0); + + ggml_tensor * lid_source_kv = ggml_concat(ctx0, lid_base_kv, lid_state_kv, 1); + ggml_tensor * lid_source_score = ggml_concat(ctx0, lid_base_score, lid_state_score, 1); ggml_tensor * kv_comp_lid_state = build_overlap_compressed_kv_from_state( lid_source_kv, @@ -982,8 +1112,19 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( ggml_build_forward_expand(gf, inp_dsv4->mctx->get_lid()->cpy_k(ctx0, kv_comp_lid_state, inp_dsv4->get_lid().state_write_idxs, il)); - lid_state_kv = dsv4_with_zero_dep(ctx0, lid_state_kv, kv_comp_lid_state); - lid_state_score = dsv4_with_zero_dep(ctx0, lid_state_score, kv_comp_lid_state); + ggml_tensor * lid_snapshot_source_kv = ggml_concat(ctx0, + lid_restored.kv, lid_state_kv, 1); + ggml_tensor * lid_snapshot_source_score = ggml_concat(ctx0, + lid_restored.score, lid_state_score, 1); + + const dsv4_state_tensors lid_snapshot = dsv4_build_state_snapshot( + ctx0, inp_dsv4->get_lid(), lid_state, lid_snapshot_source_kv, lid_snapshot_source_score, il); + if (lid_snapshot.kv != nullptr) { + ggml_build_forward_expand(gf, lid_snapshot.kv); + } + if (lid_snapshot.score != nullptr) { + ggml_build_forward_expand(gf, lid_snapshot.score); + } ggml_tensor * lid_persist_kv = ggml_get_rows(ctx0, lid_state_kv, inp_dsv4->get_lid().state_persist_src_idxs); ggml_tensor * lid_persist_score = ggml_get_rows(ctx0, lid_state_score, inp_dsv4->get_lid().state_persist_src_idxs); @@ -997,15 +1138,21 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( ggml_build_forward_expand(gf, lid_state_score); } - ggml_tensor * hca_state_dep = nullptr; + const llama_dsv4_comp_state * hca_state = nullptr; + dsv4_state_tensors hca_restored = {}; if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_write_idxs) { GGML_ASSERT(hca_state_kv); GGML_ASSERT(hca_state_score); - ggml_tensor * hca_source_kv = ggml_concat(ctx0, - inp_dsv4->mctx->get_hca_state()->get_kv(ctx0, il), hca_state_kv, 1); - ggml_tensor * hca_source_score = ggml_concat(ctx0, - inp_dsv4->mctx->get_hca_state()->get_score(ctx0, il), hca_state_score, 1); + hca_state = inp_dsv4->mctx->get_hca_state(); + hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il); + ggml_tensor * hca_base_kv = dsv4_view_2d( + ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0); + ggml_tensor * hca_base_score = dsv4_view_2d( + ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0); + + hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1); + hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1); ggml_tensor * kv_comp_hca = build_hca_compressed_kv_from_state( hca_source_kv, @@ -1024,15 +1171,41 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( ggml_build_forward_expand(gf, inp_dsv4->mctx->get_hca()->cpy_k(ctx0, kv_comp_hca, inp_dsv4->get_hca().state_write_idxs, il)); - hca_state_dep = kv_comp_hca; } if (ratio == DSV4_HCA_RATIO && inp_dsv4->get_hca().state_pos) { GGML_ASSERT(hca_state_kv); GGML_ASSERT(hca_state_score); - hca_state_kv = dsv4_with_zero_dep(ctx0, hca_state_kv, hca_state_dep); - hca_state_score = dsv4_with_zero_dep(ctx0, hca_state_score, hca_state_dep); + if (hca_state == nullptr) { + hca_state = inp_dsv4->mctx->get_hca_state(); + } + if (hca_restored.kv == nullptr) { + hca_restored = dsv4_build_state_restore(ctx0, inp_dsv4->get_hca(), hca_state, il); + } + if (hca_source_kv == nullptr || hca_source_score == nullptr) { + ggml_tensor * hca_base_kv = dsv4_view_2d( + ctx0, hca_restored.kv, hca_restored.kv->ne[0], hca_state->get_n_rows(), 0); + ggml_tensor * hca_base_score = dsv4_view_2d( + ctx0, hca_restored.score, hca_restored.score->ne[0], hca_state->get_n_rows(), 0); + + hca_source_kv = ggml_concat(ctx0, hca_base_kv, hca_state_kv, 1); + hca_source_score = ggml_concat(ctx0, hca_base_score, hca_state_score, 1); + } + + ggml_tensor * hca_snapshot_source_kv = ggml_concat(ctx0, + hca_restored.kv, hca_state_kv, 1); + ggml_tensor * hca_snapshot_source_score = ggml_concat(ctx0, + hca_restored.score, hca_state_score, 1); + + const dsv4_state_tensors hca_snapshot = dsv4_build_state_snapshot( + ctx0, inp_dsv4->get_hca(), hca_state, hca_snapshot_source_kv, hca_snapshot_source_score, il); + if (hca_snapshot.kv != nullptr) { + ggml_build_forward_expand(gf, hca_snapshot.kv); + } + if (hca_snapshot.score != nullptr) { + ggml_build_forward_expand(gf, hca_snapshot.score); + } ggml_tensor * hca_persist_kv = ggml_get_rows(ctx0, hca_state_kv, inp_dsv4->get_hca().state_persist_src_idxs); ggml_tensor * hca_persist_score = ggml_get_rows(ctx0, hca_state_score, inp_dsv4->get_hca().state_persist_src_idxs); @@ -1047,7 +1220,14 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention( } ggml_tensor * out = nullptr; - if (ratio == DSV4_CSA_RATIO && + if (inp_mtp) { + out = build_attn(inp_mtp, + nullptr, nullptr, nullptr, + q, kv, nullptr, + nullptr, layer.attn_sinks, nullptr, + 1.0f/sqrtf(float(n_embd_head)), il); + cb(out, "attn_raw", il); + } else if (ratio == DSV4_CSA_RATIO && inp_dsv4->get_csa().kq_mask && inp_dsv4->get_lid().kq_mask && inp_dsv4->get_lid().k_rot) { @@ -1106,6 +1286,12 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p cb(inpL, "hc_init", -1); for (int il = 0; il < n_layer; ++il) { + if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) { + res->t_layer_inp[il] = dsv4_hc_mean(ctx0, inpL); + cb(res->t_layer_inp[il], "layer_inp", il); + ggml_build_forward_expand(gf, res->t_layer_inp[il]); + } + ggml_tensor * residual = inpL; ggml_tensor * post = nullptr; ggml_tensor * comb = nullptr; @@ -1182,10 +1368,23 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p cb(inpL, "l_last", il); } + if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) { + res->t_layer_inp[n_layer] = dsv4_hc_mean(ctx0, inpL); + cb(res->t_layer_inp[n_layer], "layer_inp", n_layer); + ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); + } + + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + ggml_tensor * flat_out = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat; + + if (cparams.embeddings_nextn) { + ggml_tensor * h_nextn = cparams.embeddings_nextn_masked ? flat_out : inpL; + cb(h_nextn, "h_nextn", -1); + res->t_h_nextn = h_nextn; + } + if (inp_out_ids) { - ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); - flat = ggml_get_rows(ctx0, flat, inp_out_ids); - inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); + inpL = ggml_reshape_3d(ctx0, flat_out, n_embd, hc, n_outputs); } cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base); @@ -1201,3 +1400,145 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p ggml_build_forward_expand(gf, cur); } + + +llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : + graph(params) { + GGML_ASSERT(hparams.n_layer_nextn > 0 && "DEEPSEEK4 MTP requires n_layer_nextn > 0"); + GGML_ASSERT(hparams.n_layer_nextn == 1 && "DEEPSEEK4 MTP currently only supports a single MTP block"); + GGML_ASSERT(cparams.nextn_layer_offset >= 0 && + cparams.nextn_layer_offset < (int) hparams.n_layer_nextn && + "nextn_layer_offset out of range [0, n_layer_nextn)"); + GGML_ASSERT(ubatch.token && "DEEPSEEK4 MTP requires token input"); + + const int64_t hc = hparams.dsv4_hc_mult; + GGML_ASSERT(hparams.n_embd_out() == (uint32_t) (n_embd*hc) && "DEEPSEEK4 MTP hidden width mismatch"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + + auto inp = std::make_unique(hparams.n_embd_out()); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens); + ggml_set_input(inp->embd); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + ggml_tensor * tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + cb(tok_embd, "mtp_tok_embd", il); + + ggml_tensor * h_state = ggml_reshape_3d(ctx0, inp->h, n_embd, hc, n_tokens); + cb(h_state, "mtp_h_state", il); + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa(); + + ggml_tensor * h_norm = build_norm(h_state, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + e_norm = ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens); + e_norm = ggml_repeat_4d(ctx0, e_norm, n_embd, hc, n_tokens, 1); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * inpL = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(inpL, "mtp_eh_proj", il); + + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + + ggml_tensor * cur = build_hc_pre(inpL, + layer.hc_attn_fn, + layer.hc_attn_scale, + layer.hc_attn_base, + &post, &comb, il); + cb(cur, "mtp_hc_attn_pre", il); + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + cur = build_attention(model, inp_attn, cur, inp_pos, il); + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "mtp_hc_attn_post", il); + + residual = inpL; + cur = build_hc_pre(inpL, + layer.hc_ffn_fn, + layer.hc_ffn_scale, + layer.hc_ffn_base, + &post, &comb, il); + cb(cur, "mtp_hc_ffn_pre", il); + + cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_ffn_norm", il); + + GGML_ASSERT((uint32_t) il >= hparams.dsv4_hash_layer_count && "DEEPSEEK4 MTP does not support hash-routed MTP blocks"); + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, hparams.n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(moe_out, "mtp_ffn_moe_out", il); + + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "mtp_ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "mtp_ffn_out", il); + + inpL = build_hc_post(cur, residual, post, comb, il); + inpL = build_cvec(inpL, il); + cb(inpL, "mtp_l_out", il); + + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + ggml_tensor * h_nextn = ggml_get_rows(ctx0, flat, inp_out_ids); + cb(h_nextn, "h_nextn", -1); + res->t_h_nextn = h_nextn; + + inpL = ggml_reshape_3d(ctx0, h_nextn, n_embd, hc, n_outputs); + + cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base); + cb(cur, "mtp_hc_head", -1); + + ggml_tensor * head_norm_w = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm; + GGML_ASSERT(head_norm_w && "DEEPSEEK4 MTP missing shared head norm"); + cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); + cb(cur, "mtp_shared_head_norm", -1); + res->t_embd = cur; + + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + GGML_ASSERT(head_w && "DEEPSEEK4 MTP missing LM head"); + cur = ggml_mul_mat(ctx0, head_w, cur); + cb(cur, "result_output", -1); + + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index add0bd299f8..0a8809e250f 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -20,6 +20,48 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { } LLAMA_LOG_INFO("]\n"); + // DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring) + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false); + if (hparams.dsv4_hc_mult > 0) { + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all); + if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) { + hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp; + } + ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count); + ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank); + ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); + ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false); + + if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { + throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring"); + } + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + if (hparams.dsv4_compress_ratios[il] != 0) { + throw std::runtime_error("DSpark DSV4 draft expects uncompressed attention on all stages"); + } + } + + GGML_ASSERT(hparams.n_swa > 0); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + hparams.set_swa_pattern(0); + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + hparams.is_swa_impl[il] = true; + } + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + + type = LLM_TYPE_UNKNOWN; + return; + } + // optional interleaved sliding-window attention with per-layer pattern array. // DFlash has a single rope, so the SWA rope == main rope. if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) { @@ -69,6 +111,57 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader & ml) { aux_norm = create_tensor(tn(LLM_TENSOR_ENC_AUX_NORM, "weight"), { n_embd, (int64_t) target_layer_ids.size() }, 0); } + if (hparams.dsv4_hc_mult > 0) { + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_expert_shared = hparams.n_expert_shared; + const int64_t n_embd_head = hparams.n_embd_head_k(); + const int64_t o_groups = hparams.dsv4_o_group_count; + const int64_t o_lora_rank = hparams.dsv4_o_lora_rank; + const int64_t hc_mult = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc_mult * n_embd; + const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult; + + hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0); + hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0); + hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0); + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0); + layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0); + layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0); + layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank * o_groups}, 0); + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0); + + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0); + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0); + } + return; + + } + for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; @@ -111,6 +204,9 @@ std::unique_ptr llama_model_dflash::build_arch_graph(const ll return std::make_unique>(*this, params); case LLM_GRAPH_TYPE_DEFAULT: case LLM_GRAPH_TYPE_DECODER: + if (hparams.dsv4_hc_mult > 0) { + return std::make_unique(*this, params); + } return std::make_unique>(*this, params); default: GGML_ABORT("invalid graph type"); @@ -482,3 +578,178 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra build_dspark_markov_head(*this, model, inp_tokens); } } + +// DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above): +// * embd batch -> project main_x through each stage's wkv and inject K into the ring cache +// * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads +llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_graph_params & params) : + llama_model_deepseek4::graph(params) { + const int64_t n_embd_head = hparams.n_embd_head_k(); + const int64_t n_embd_head_rope = hparams.n_rot(); + const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope; + + ggml_tensor * inp_pos = build_inp_pos(); + + llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa(); + + // KV cache injection: fused target features from the encoder + if (ubatch.embd) { + auto inp = std::make_unique(n_embd); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * inp_g = inp->embd; + cb(inp_g, "inp_g_embeddings", -1); + + res->add_input(std::move(inp)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + // main-track KV: kv_norm(wkv(main_x)) with rope on the trailing dims, same + // rope parameters as the uncompressed layers in build_attention_impl + ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g); + kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il); + kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens); + + ggml_tensor * kv_nope = ggml_view_3d(ctx0, kv, n_embd_head_nope, 1, n_tokens, + ggml_row_size(kv->type, n_embd_head), + ggml_row_size(kv->type, n_embd_head), + 0); + ggml_tensor * kv_pe = ggml_view_3d(ctx0, kv, n_embd_head_rope, 1, n_tokens, + ggml_row_size(kv->type, n_embd_head), + ggml_row_size(kv->type, n_embd_head), + ggml_row_size(kv->type, n_embd_head_nope)); + kv_pe = ggml_rope_ext(ctx0, kv_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, 0, + freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + kv = ggml_concat(ctx0, kv_nope, kv_pe, 0); + cb(kv, "kv_injected", il); + + if (inp_attn->self_k_rot_swa) { + kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa); + } + ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il)); + } + + res->t_embd = inp_g; + + ggml_build_forward_expand(gf, inp_g); + return; + } + + // tok_embd from the target model (shared via ctx_other) + auto * tok_embd = model.tok_embd; + if (tok_embd == nullptr) { + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + + GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings"); + tok_embd = model_other->tok_embd; + } + + auto inp = std::make_unique(n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + ggml_tensor * inp_tokens = inp->tokens; + + ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); + cb(inpL, "inp_noise_embd", -1); + + res->add_input(std::move(inp)); + + const int64_t hc = hparams.dsv4_hc_mult; + inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens); + inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); + cb(inpL, "hc_init", -1); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + + ggml_tensor * cur = build_hc_pre(inpL, + layer.hc_attn_fn, + layer.hc_attn_scale, + layer.hc_attn_base, + &post, &comb, il); + cb(cur, "hc_attn_pre", il); + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + cur = build_attention(model, inp_attn, cur, inp_pos, il); + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + residual = inpL; + cur = build_hc_pre(inpL, + layer.hc_ffn_fn, + layer.hc_ffn_scale, + layer.hc_ffn_base, + &post, &comb, il); + cb(cur, "hc_ffn_pre", il); + + cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, hparams.n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "l_out", il); + } + + ggml_tensor * cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base); + cb(cur, "hc_head", -1); + + // confidence head input: the reference scores the pre-norm collapsed hidden state + res->t_embd = cur; + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + + // lm_head from the target model (shared via ctx_other) + auto * output = model.output; + if (output == nullptr) { + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection"); + output = model_other->output; + } + + cur = build_lora_mm(output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); + + if (model.dspark_markov_w1) { + build_dspark_markov_head(*this, model, inp_tokens); + } +} diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp index 4080a934cb9..d50e186cce9 100644 --- a/src/models/mimo2.cpp +++ b/src/models/mimo2.cpp @@ -30,7 +30,11 @@ void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) { const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); - const int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + + if (!ml.load_mtp) { + mtp_flags |= TENSOR_SKIP; + } tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); diff --git a/src/models/models.h b/src/models/models.h index 0a9e7e0707b..1051f9fa030 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1106,6 +1106,7 @@ struct llama_model_deepseek4 : public llama_model_base { void load_arch_tensors(llama_model_loader & ml) override; struct graph : public llm_graph_context { + graph(const llm_graph_params & params) : llm_graph_context(params) {} graph(const llama_model & model, const llm_graph_params & params); ggml_tensor * build_hc_pre( @@ -1137,6 +1138,21 @@ struct llama_model_deepseek4 : public llama_model_base { ggml_tensor * inp_pos, int il) const; + ggml_tensor * build_attention( + const llama_model & model, + llm_graph_input_attn_k_iswa * inp_mtp, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int il) const; + + ggml_tensor * build_attention_impl( + const llama_model & model, + llm_graph_input_dsv4 * inp_dsv4, + llm_graph_input_attn_k_iswa * inp_mtp, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int il) const; + ggml_tensor * build_hca_compressed_kv_from_state( ggml_tensor * kv_state, ggml_tensor * score_state, @@ -1212,6 +1228,10 @@ struct llama_model_deepseek4 : public llama_model_base { int il) const; }; + struct graph_mtp : public graph { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; @@ -1280,6 +1300,10 @@ struct llama_model_dflash : public llama_model_base { ggml_tensor * build_inp_embd_enc() const; }; + struct graph_dsv4 : public llama_model_deepseek4::graph { + graph_dsv4(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr 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 d84c508c2d6..a59af37338f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8332,6 +8332,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_comb(1, 1)); test_cases.emplace_back(new test_dsv4_hc_comb(17, 4)); test_cases.emplace_back(new test_dsv4_hc_comb(257, 8)); + test_cases.emplace_back(new test_dsv4_hc_comb(17, 20)); test_cases.emplace_back(new test_dsv4_hc_pre(1, 1)); test_cases.emplace_back(new test_dsv4_hc_pre(31, 17)); @@ -8341,6 +8342,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_dsv4_hc_post(1, 1)); test_cases.emplace_back(new test_dsv4_hc_post(31, 17)); test_cases.emplace_back(new test_dsv4_hc_post(128, 257)); + test_cases.emplace_back(new test_dsv4_hc_post(4096, 21)); // glu ops for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { @@ -8830,6 +8832,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_repeat(GGML_TYPE_F32, {10, 5, 4, ne3}, {1, 2, 1, 1})); test_cases.emplace_back(new test_repeat(GGML_TYPE_F32, {10, 5, 4, ne3}, {1, 1, 2, 1})); test_cases.emplace_back(new test_repeat(GGML_TYPE_F32, {10, 5, 4, ne3}, {1, 1, 1, 2})); + test_cases.emplace_back(new test_repeat(GGML_TYPE_F16, {10, 5, 4, ne3}, {2, 1, 1, 1})); test_cases.emplace_back(new test_repeat(GGML_TYPE_I32, {10, 5, 4, ne3}, {2, 1, 1, 1})); test_cases.emplace_back(new test_repeat(GGML_TYPE_I16, {10, 5, 4, ne3}, {1, 1, 1, 2})); test_cases.emplace_back(new test_repeat(GGML_TYPE_BF16, {10, 5, 4, ne3}, {2, 1, 1, 1})); diff --git a/tests/test-chat-peg-parser.cpp b/tests/test-chat-peg-parser.cpp index 908b13fd0ca..3ab7a67b6a8 100644 --- a/tests/test-chat-peg-parser.cpp +++ b/tests/test-chat-peg-parser.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include "nlohmann/json.hpp" @@ -21,6 +22,7 @@ static void test_example_qwen3_non_coder(testing & t); static void test_command7_parser_compare(testing & t); static void test_prefix_tool_names(testing & t); static void test_tagged_peg_parser(testing & t); +static void test_permute(testing & t); int main(int argc, char * argv[]) { testing t(std::cout); @@ -39,6 +41,7 @@ int main(int argc, char * argv[]) { t.test("comparison", test_command7_parser_compare); t.test("prefix tool names", test_prefix_tool_names); t.test("tagged peg parser", test_tagged_peg_parser); + t.test("permute", test_permute); return t.summary(); } @@ -981,3 +984,103 @@ static void test_tagged_peg_parser(testing & t) { t.assert_equal("fun_post should be '>'", ">", result.tags["fun_post"]); }); } + +static void test_permute(testing & t) { + auto accepts = [](const common_peg_arena & parser, const std::string & input) { + common_peg_parse_context ctx(input); + return parser.parse(ctx).success(); + }; + + auto gbnf_of = [](const common_peg_arena & parser) { + return build_grammar([&](const common_grammar_builder & builder) { parser.build_grammar(builder); }); + }; + + auto assert_gbnf_equal = [](testing & t, const std::string & expected, const std::string & actual) { + static const std::regex leading_ws_re = std::regex(R"((^|\n)\s+)"); + t.assert_equal("gbnf are equal", std::regex_replace(expected, leading_ws_re, "$1"), actual); + }; + + auto count_rules = [](const std::string & gbnf, const std::string & prefix) { + size_t count = 0; + for (const auto & line : string_split(gbnf, '\n')) { + if (line.rfind(prefix, 0) == 0) { + count++; + } + } + return count; + }; + + t.test("accepts every ordering", [&](testing & t) { + auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) { + return p.permute("abc", { p.literal("a"), p.literal("b"), p.literal("c") }) + p.end(); + }); + + for (const std::string input : { "abc", "acb", "bac", "bca", "cab", "cba" }) { + t.assert_true("accepts " + input, accepts(parser, input)); + } + }); + + t.test("single element", [&](testing & t) { + auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) { + return p.permute("a", { p.literal("a") }) + p.end(); + }); + + t.assert_true("accepts a", accepts(parser, "a")); + t.assert_true("rejects aa", !accepts(parser, "aa")); + }); + + t.test("grammar left-factorizes shared tails", [&](testing & t) { + auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) { + return p.permute("abc", { p.literal("a"), p.literal("b"), p.literal("c") }) + p.end(); + }); + + // Every rule is one remaining subset, keyed by bitmask: abc-3 is {a,b}, abc-7 is {a,b,c}. + // Each subset is emitted once and shared by every branch that leads into it. + assert_gbnf_equal(t, R"""( + abc-1 ::= "a" + abc-2 ::= "b" + abc-3 ::= "a" abc-2 | "b" abc-1 + abc-4 ::= "c" + abc-5 ::= "a" abc-4 | "c" abc-1 + abc-6 ::= "b" abc-4 | "c" abc-2 + abc-7 ::= "a" abc-6 | "b" abc-5 | "c" abc-3 + root ::= abc-7 + space ::= | " " | "\n"{1,2} [ \t]{0,20} + )""", gbnf_of(parser)); + }); + + t.test("grammar emits one rule per remaining subset", [&](testing & t) { + auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) { + return p.permute("abcd", { p.literal("a"), p.literal("b"), p.literal("c"), p.literal("d") }) + p.end(); + }); + + // 2^4 - 1 non-empty subsets, one rule each - not the 4! = 24 orderings. + t.assert_equal("permute rule count", 15u, count_rules(gbnf_of(parser), "abcd-")); + }); + + t.test("grammar emits no rules for a single element", [&](testing & t) { + auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) { + return p.permute("a", { p.literal("a") }) + p.end(); + }); + + assert_gbnf_equal(t, R"""( + root ::= "a" + space ::= | " " | "\n"{1,2} [ \t]{0,20} + )""", gbnf_of(parser)); + }); + + t.test("grammar falls back to the given order when too large", [&](testing & t) { + auto parser = build_chat_peg_parser([](common_chat_peg_builder & p) { + std::vector parsers; + for (size_t i = 0; i <= COMMON_CHAT_MAX_PERMUTE; i++) { + parsers.push_back(p.literal(std::string(1, (char) ('a' + i)))); + } + return p.permute("big", parsers) + p.end(); + }); + + assert_gbnf_equal(t, R"""( + root ::= "a" "b" "c" "d" "e" "f" "g" + space ::= | " " | "\n"{1,2} [ \t]{0,20} + )""", gbnf_of(parser)); + }); +} diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index ef02fdde57e..8d571d369e5 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -2278,46 +2278,39 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})") .run(); - // tool call segment in reasoning + // a tool call ends the prefilled thinking block, with or without a closing tst.test( - "Let's call a tool: \n" - "\n" - "\n" - "def hello():\n" - " print(\"Not the real call!\")\n" - "\n" - "hello()\n" + "\n" + "\n" + "\n" + "pwd\n" "\n" "\n" - "\n\n\n" + "") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ run_in_terminal_tool }) + .expect_tool_calls({ + { "run_in_terminal", R"({"command": "pwd"})", {} }, + }) + .run(); + + // ...including after the model has thought about it + tst.test( + "Need to inspect the current directory.\n" "\n" - "\n" - "\n" - "def hello():\n" - " print(\"Hello, world!\")\n" - "\n" - "hello()\n" + "\n" + "\n" + "pwd\n" "\n" "\n" "") .enable_thinking(true) .reasoning_format(COMMON_REASONING_FORMAT_AUTO) - .tools({ - python_tool - }) - .expect_reasoning( - "Let's call a tool: \n" - "\n" - "\n" - "def hello():\n" - " print(\"Not the real call!\")\n" - "\n" - "hello()\n" - "\n" - "\n" - "") + .tools({ run_in_terminal_tool }) + .expect_reasoning("Need to inspect the current directory.") .expect_tool_calls({ - { "python", "{\"code\": \"def hello():\\n print(\\\"Hello, world!\\\")\\n\\nhello()\"}", {} }, + { "run_in_terminal", R"({"command": "pwd"})", {} }, }) .run(); @@ -2461,17 +2454,6 @@ static void test_template_output_peg_parsers(bool detailed_debug) { }) .run(); - tst.test( - "I might call later, but I am still thinking.\n" - "\n\n" - "Final answer without tools.") - .reasoning_format(COMMON_REASONING_FORMAT_AUTO) - .enable_thinking(true) - .tools({ run_in_terminal_tool }) - .expect_reasoning("I might call later, but I am still thinking.") - .expect_content("Final answer without tools.") - .run(); - // Continuation tests tst.test("world!\nWhat's up?") .reasoning_format(COMMON_REASONING_FORMAT_AUTO) @@ -2776,49 +2758,6 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})") .run(); - // tool call segment in reasoning - tst.test( - "Let's call a tool: \n" - "\n" - "\n" - "def hello():\n" - " print(\"Not the real call!\")\n" - "\n" - "hello()\n" - "\n" - "\n" - "\n\n" - "\n" - "\n" - "\n" - "def hello():\n" - " print(\"Hello, world!\")\n" - "\n" - "hello()\n" - "\n" - "\n" - "\n" - ) - .enable_thinking(true) - .reasoning_format(COMMON_REASONING_FORMAT_AUTO) - .tools({ - python_tool - }) - .expect_reasoning("Let's call a tool: \n" - "\n" - "\n" - "def hello():\n" - " print(\"Not the real call!\")\n" - "\n" - "hello()\n" - "\n" - "\n" - "\n") - .expect_tool_calls({ - { "python", "{\"code\": \"def hello():\\n print(\\\"Hello, world!\\\")\\n\\nhello()\"}", {} }, - }) - .run(); - // Continuation tests tst.test("world!\nWhat's up?") .reasoning_format(COMMON_REASONING_FORMAT_AUTO) @@ -3572,6 +3511,61 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .expect_reconstruction() .run(); + // Some models skip the opening and go straight to + tst.test( + "\n" + "\n" + "1\n" + "\n" + "\n" + "") + .tools({ special_function_tool }) + .expect(message_assist_call) + .run(); + + tst.test( + "Let me call it.\n" + "\n" + "\n" + "1\n" + "\n" + "\n" + "") + .tools({ special_function_tool }) + .expect_content("Let me call it.\n") + .expect_tool_calls({ + { "special_function", R"({"arg1": 1})", {} }, + }) + .run(); + + // Only the first call may omit it, the rest keep the \n separator + tst.test( + "\n" + "\n" + "1\n" + "\n" + "\n" + "\n" + "\n" + "\n" + "\n" + "1\n" + "\n" + "\n" + "2\n" + "\n" + "\n" + "") + .parallel_tool_calls(true) + .tools({ + special_function_tool, special_function_tool_with_optional_param + }) + .expect_tool_calls({ + { "special_function", R"({"arg1": 1})", {} }, + { "special_function_with_opt", R"({"arg1": 1, "arg2": 2})", {} }, + }) + .run(); + tst.test( "\n" "\n" @@ -3680,6 +3674,37 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .expect_reconstruction() .run(); + // Test flexible required argument ordering (required args still come first, in any order) + tst.test( + "\n" + "\n" + "\n#include\n\n" + "\nfoo.c\n\n" + "\n#iclunde\n\n" + "\n" + "") + .tools({ edit_tool }) + .expect_tool_calls({ + { "edit", R"({"newString": "#include", "filename": "foo.c", "oldString": "#iclunde"})", {} }, + }) + .expect_reconstruction() + .run(); + + tst.test( + "\n" + "\n" + "\n42\n\n" + "\nhello\n\n" + "\n200\n\n" + "\n" + "") + .tools({ tool_2req_4opt }) + .expect_tool_calls({ + { "tool_2req_4opt", R"({"req2": 42, "req1": "hello", "opt2": 200})", {} }, + }) + .expect_reconstruction() + .run(); + // Test flexible optional argument ordering (2 required + 4 optional, reversed optional order) tst.test( "\n" diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index a1ed2a76f87..4336e4e13d4 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -430,7 +430,7 @@ static bool arch_supported(const llm_arch arch) { // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MINIMAX_M3) { return false; } #endif // GGML_USE_WEBGPU diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index 8e2eace6a19..5d1f0140b62 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -83,27 +83,33 @@ int main(int argc, char ** argv) { if (llama_vocab_type(vocab) == LLAMA_VOCAB_TYPE_NONE) { tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 }; } else { - tokens = common_tokenize(ctx_src, "The quick brown fox jumps", true); + tokens = common_tokenize(ctx_src, "The quick brown fox jumps over the lazy dog", true); } const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src); - if (tokens.size() > n_rs_seq + 1) { - tokens.resize(n_rs_seq + 1); + constexpr uint32_t n_rollback = 3; + if (n_rs_seq < n_rollback) { + fprintf(stderr, "%s : skipping because n_rs_seq is too small\n", __func__); + llama_free(ctx_src); + llama_free(ctx_dst); + return 0; } - if (tokens.size() < 2) { + if (tokens.empty()) { fprintf(stderr, "%s : not enough prompt tokens\n", __func__); return 1; } - const uint32_t n_tokens = tokens.size(); - const llama_token last_tok = tokens.back(); - const llama_pos last_pos = (llama_pos) n_tokens - 2; + tokens.resize(n_rs_seq + 1, tokens.back()); - // Decode the full prompt on the source, then roll back the last position. + const uint32_t n_tokens = tokens.size(); + const llama_pos rollback_pos = (llama_pos) n_tokens - n_rollback; + + // Decode the full prompt on the source, then roll back three positions. + // Replaying them crosses DSV4's ratio-4 compressor boundary. // Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0). if (!decode_tokens(ctx_src, tokens, n_tokens)) { fprintf(stderr, "%s : failed to decode prompt\n", __func__); return 1; } - if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, last_pos, -1)) { + if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1)) { fprintf(stderr, "%s : rollback failed\n", __func__); return 1; } @@ -113,31 +119,56 @@ int main(int argc, char ** argv) { ckpt.update_tgt(ctx_src, 0, 0); ckpt.load_tgt(ctx_dst, 0, 0); - // Replay the rolled-back token on both contexts and compare logits. - if (!decode_one(ctx_src, last_tok, last_pos) || - !decode_one(ctx_dst, last_tok, last_pos)) { - fprintf(stderr, "%s : replay failed\n", __func__); + constexpr float eps = 1e-5f; + std::vector> logits_src_replay(n_rollback); + const auto replay_and_compare = [&](const char * mode) { + for (uint32_t i = 0; i < n_rollback; ++i) { + const llama_pos pos = rollback_pos + i; + if (!decode_one(ctx_src, tokens[pos], pos) || + !decode_one(ctx_dst, tokens[pos], pos)) { + fprintf(stderr, "%s : %s replay failed at position %d\n", __func__, mode, pos); + return false; + } + + const float * logits_src = llama_get_logits_ith(ctx_src, 0); + const float * logits_dst = llama_get_logits_ith(ctx_dst, 0); + if (logits_src == nullptr || logits_dst == nullptr) { + fprintf(stderr, "%s : missing %s logits at position %d\n", __func__, mode, pos); + return false; + } + + logits_src_replay[i].assign(logits_src, logits_src + n_vocab); + for (int token = 0; token < n_vocab; ++token) { + if (std::fabs(logits_src[token] - logits_dst[token]) > eps) { + fprintf(stderr, "%s : %s logits mismatch at position %d, token %d (%g != %g)\n", + __func__, mode, pos, token, (double) logits_src[token], (double) logits_dst[token]); + return false; + } + } + } + return true; + }; + if (!replay_and_compare("full")) { return 1; } - const float * logits_src = llama_get_logits_ith(ctx_src, 0); - const float * logits_dst = llama_get_logits_ith(ctx_dst, 0); - if (logits_src == nullptr || logits_dst == nullptr) { - fprintf(stderr, "%s : missing logits\n", __func__); + if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) || + !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) { + fprintf(stderr, "%s : partial rollback failed\n", __func__); return 1; } - constexpr float eps = 1e-5f; - for (int i = 0; i < n_vocab; ++i) { - if (std::fabs(logits_src[i] - logits_dst[i]) > eps) { - fprintf(stderr, "%s : logits mismatch at token %d (%g != %g)\n", - __func__, i, (double) logits_src[i], (double) logits_dst[i]); - return 1; - } + constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; + common_prompt_checkpoint ckpt_partial; + ckpt_partial.update_tgt(ctx_src, 0, partial_flags); + ckpt_partial.load_tgt(ctx_dst, 0, partial_flags); + + if (!replay_and_compare("partial")) { + return 1; } // Repeat the load into a context that already has its own rollback state: - // groups 1..n_rs_seq hold a *different* prompt's history, and rs_idx[0] is + // groups 1..n_rs_seq hold a different prompt's history, and rs_idx[0] is // non-zero at load time. The restore must wipe that state and still match. llama_context * ctx_dirty = make_ctx(params, model); if (ctx_dirty == nullptr) { @@ -156,30 +187,33 @@ int main(int argc, char ** argv) { fprintf(stderr, "%s : dirty prompt decode failed\n", __func__); return 1; } - if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, last_pos, -1)) { + if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, rollback_pos, -1)) { fprintf(stderr, "%s : dirty rollback failed\n", __func__); return 1; } ckpt.load_tgt(ctx_dirty, 0, 0); - if (!decode_one(ctx_dirty, last_tok, last_pos)) { - fprintf(stderr, "%s : dirty replay failed\n", __func__); - return 1; - } - - const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0); - if (logits_dirty == nullptr) { - fprintf(stderr, "%s : missing dirty logits\n", __func__); - return 1; - } + for (uint32_t i = 0; i < n_rollback; ++i) { + const llama_pos pos = rollback_pos + i; + if (!decode_one(ctx_dirty, tokens[pos], pos)) { + fprintf(stderr, "%s : dirty replay failed at position %d\n", __func__, pos); + return 1; + } - for (int i = 0; i < n_vocab; ++i) { - if (std::fabs(logits_src[i] - logits_dirty[i]) > eps) { - fprintf(stderr, "%s : dirty-ctx logits mismatch at token %d (%g != %g)\n", - __func__, i, (double) logits_src[i], (double) logits_dirty[i]); + const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0); + if (logits_dirty == nullptr) { + fprintf(stderr, "%s : missing dirty logits at position %d\n", __func__, pos); return 1; } + + for (int token = 0; token < n_vocab; ++token) { + if (std::fabs(logits_src_replay[i][token] - logits_dirty[token]) > eps) { + fprintf(stderr, "%s : dirty-ctx logits mismatch at position %d, token %d (%g != %g)\n", + __func__, pos, token, (double) logits_src_replay[i][token], (double) logits_dirty[token]); + return 1; + } + } } fprintf(stderr, "%s : recurrent rollback checkpoint restored successfully\n", __func__); diff --git a/tools/ui/CMakeLists.txt b/tools/ui/CMakeLists.txt index 74bca417e37..208b46a5c15 100644 --- a/tools/ui/CMakeLists.txt +++ b/tools/ui/CMakeLists.txt @@ -61,12 +61,30 @@ if(CMAKE_CROSSCOMPILING) # phony target to tie it into the dependency graph add_custom_target(llama-ui-embed DEPENDS "${LLAMA_UI_EMBED_EXE}") else() + # exclude llama-ui-embed from sanitizer flags, + # it's a build-time-only tool, no need to instrument it + # this is to fix TSan "memory layout is incompatible" error on CI + get_directory_property(_llama_ui_dir_co COMPILE_OPTIONS) + get_directory_property(_llama_ui_dir_ll LINK_LIBRARIES) + set(_llama_ui_embed_co ${_llama_ui_dir_co}) + set(_llama_ui_embed_ll ${_llama_ui_dir_ll}) + list(FILTER _llama_ui_embed_co EXCLUDE REGEX ".*-fsanitize=.*") + list(FILTER _llama_ui_embed_ll EXCLUDE REGEX ".*-fsanitize=.*") + set_directory_properties(PROPERTIES + COMPILE_OPTIONS "${_llama_ui_embed_co}" + LINK_LIBRARIES "${_llama_ui_embed_ll}") + add_executable(llama-ui-embed embed.cpp) target_compile_features(llama-ui-embed PRIVATE cxx_std_17) set_target_properties(llama-ui-embed PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" ) set(LLAMA_UI_EMBED_EXE "$") + + # restore so the llama-ui library below keeps sanitizer instrumentation + set_directory_properties(PROPERTIES + COMPILE_OPTIONS "${_llama_ui_dir_co}" + LINK_LIBRARIES "${_llama_ui_dir_ll}") endif() # Run the provisioning script every build so source changes in tools/ui/ are diff --git a/vendor/cpp-httplib/CMakeLists.txt b/vendor/cpp-httplib/CMakeLists.txt index c6eb372b5f2..bf674583ff9 100644 --- a/vendor/cpp-httplib/CMakeLists.txt +++ b/vendor/cpp-httplib/CMakeLists.txt @@ -41,7 +41,7 @@ if (LLAMA_BUILD_BORINGSSL) set(FIPS OFF CACHE BOOL "Enable FIPS (BoringSSL)") set(BORINGSSL_GIT "https://boringssl.googlesource.com/boringssl" CACHE STRING "BoringSSL git repository") - set(BORINGSSL_VERSION "0.20260728.0" CACHE STRING "BoringSSL version") + set(BORINGSSL_VERSION "0.20260730.0" CACHE STRING "BoringSSL version") message(STATUS "Fetching BoringSSL version ${BORINGSSL_VERSION}")