diff --git "a/llama.cpp-openpangu-v0.7.0.diff" "b/llama.cpp-openpangu-v0.7.0.diff" new file mode 100644--- /dev/null +++ "b/llama.cpp-openpangu-v0.7.0.diff" @@ -0,0 +1,3624 @@ +diff --git a/common/arg.cpp b/common/arg.cpp +index 8b302d5d7..49354d98b 100644 +--- a/common/arg.cpp ++++ b/common/arg.cpp +@@ -3096,11 +3096,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex + ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("HF_TOKEN")); + add_opt(common_arg( + {"--mtp"}, +- "also download the multi-token prediction (MTP) head, if available (default: unused)", ++ "enable multi-token prediction (MTP) self-speculative decoding, and download the MTP head if available (default: unused)", + [](common_params & params) { + params.speculative.types.push_back(COMMON_SPECULATIVE_TYPE_DRAFT_MTP); + } +- ).set_examples({LLAMA_EXAMPLE_DOWNLOAD})); ++ ).set_examples({LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--dflash"}, + "also download the DFlash sidecar, if available (default: unused)", +diff --git a/common/chat.cpp b/common/chat.cpp +index 485671308..ddf9e3ac9 100644 +--- a/common/chat.cpp ++++ b/common/chat.cpp +@@ -906,9 +906,10 @@ common_chat_prompt common_chat_template_direct_apply_impl( + inp[k] = v; + } + } +- if (inputs.add_generation_prompt) { +- inp["add_generation_prompt"] = true; +- } ++ // always set explicitly: templates that default an undefined add_generation_prompt ++ // to true (e.g. openPangu) would otherwise render identical output for both the ++ // with- and without-generation-prompt passes, yielding an empty generation prompt ++ inp["add_generation_prompt"] = inputs.add_generation_prompt; + if (inp.contains("preserve_reasoning") && inp["preserve_reasoning"].is_boolean()) { + bool enabled = inp["preserve_reasoning"].get(); + jinja::caps_apply_preserve_reasoning(ctx, enabled); +@@ -1163,6 +1164,12 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const + return common_chat_params_init_cohere2moe; + } + ++ // openPangu - <|message_start|> turns, <|tool_call_start|> JSON tool calls and thinking reasoning ++ if (is_openpangu_template(src)) { ++ LOG_DBG("Using specialized template: openPangu\n"); ++ return common_chat_params_init_openpangu; ++ } ++ + if (is_lfm2_template(src)) { + LOG_DBG("Using specialized template: LFM2\n"); + return common_chat_params_init_lfm2_tokens; +diff --git a/common/parsers/openpangu.cpp b/common/parsers/openpangu.cpp +new file mode 100644 +index 000000000..adc5f50ce +--- /dev/null ++++ b/common/parsers/openpangu.cpp +@@ -0,0 +1,96 @@ ++#include "parsers.h" ++ ++// openPangu format detection: <|message_start|>{role} turns with ++// <|tool_call_start|>[{"name":..,"arguments":..}]<|tool_call_end|> tool calls and thinking reasoning ++bool is_openpangu_template(const std::string & src) { ++ return src.find("<|tool_call_start|>") != std::string::npos && ++ src.find("<|message_start|>") != std::string::npos; ++} ++ ++common_chat_params common_chat_params_init_openpangu(const common_chat_template & tmpl, ++ const autoparser::generation_params & inputs) { ++ common_chat_params data; ++ ++ const std::string TOOL_CALL_START = "<|tool_call_start|>"; ++ const std::string TOOL_CALL_END = "<|tool_call_end|>"; ++ const std::string THINK_START = " thinking"; ++ const std::string THINK_END = ""; ++ const std::string GEN_PROMPT = "<|message_start|>assistant\n"; ++ ++ // the template reads a 'thinking' flag (default true): thinking on force-opens the ++ // think block in the generation prompt ('...assistant\n thinking'), thinking off closes ++ // it immediately ('...assistant\n') ++ const json template_context = { {"thinking", inputs.enable_thinking} }; ++ ++ data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, std::nullopt, std::nullopt, template_context); ++ data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, std::nullopt, std::nullopt, template_context); ++ data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; ++ data.supports_thinking = true; ++ data.thinking_start_tag = THINK_START; ++ data.thinking_end_tags = { THINK_END }; ++ data.preserved_tokens = { TOOL_CALL_START, TOOL_CALL_END, THINK_START, THINK_END }; ++ ++ auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); ++ auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); ++ auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE && ++ tmpl.source().find(THINK_START) != std::string::npos; ++ 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_PROMPT + THINK_START + msg.reasoning_content; ++ if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { ++ data.generation_prompt += THINK_END + msg.render_content(); ++ } ++ ++ data.prompt.append(data.generation_prompt); ++ } ++ ++ data.parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { ++ // the parser input is generation_prompt + completion when the caller plumbs the ++ // generation prompt through (server), and just the completion otherwise (cli) -- ++ // so the assistant header is optional ++ auto generation_prompt = p.optional(p.literal(GEN_PROMPT)); ++ auto end = p.end(); ++ ++ auto reasoning = p.eps(); ++ if (extract_reasoning) { ++ // the template force-opens the think block: ' thinking' is only seen when the ++ // generation prompt is prepended, so the opening tag is optional; with thinking ++ // disabled the template emits a bare '', parsed as empty reasoning ++ reasoning = p.optional(p.optional(p.literal(THINK_START)) + p.reasoning(p.until(THINK_END)) + p.literal(THINK_END)); ++ } ++ ++ if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { ++ if (has_response_format) { ++ auto response_format = p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)); ++ return generation_prompt + reasoning + response_format + end; ++ } ++ return generation_prompt + reasoning + p.content(p.rest()) + end; ++ } ++ ++ auto tool_calls = p.standard_json_tools(TOOL_CALL_START, TOOL_CALL_END, inputs.tools, inputs.parallel_tool_calls, ++ /* force_tool_calls = */ false, ++ /* name_key = */ "name", ++ /* args_key = */ "arguments", ++ /* array_wrapped = */ true); ++ ++ auto content = p.content(p.until(TOOL_CALL_START)); ++ ++ return generation_prompt + reasoning + content + p.optional(tool_calls) + end; ++ }); ++ ++ if (include_grammar) { ++ data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); ++ data.grammar = build_grammar([&](const common_grammar_builder & builder) { ++ data.parser.build_grammar(builder, data.grammar_lazy); ++ }); ++ ++ data.grammar_triggers = { ++ { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, TOOL_CALL_START } ++ }; ++ } ++ ++ return data; ++} +diff --git a/common/parsers/parsers.h b/common/parsers/parsers.h +index 8f0a97ce3..5eaf8e8e9 100644 +--- a/common/parsers/parsers.h ++++ b/common/parsers/parsers.h +@@ -41,6 +41,8 @@ std::string common_chat_template_generation_prompt_impl( + + bool is_lfm2_template(const std::string & src); + ++bool is_openpangu_template(const std::string & src); ++ + namespace workaround { + + void convert_tool_responses_gemma4(json & messages); +@@ -80,6 +82,8 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat + + common_chat_params common_chat_params_init_muse_glimmer(const common_chat_template & tmpl, const autoparser::generation_params & inputs); + ++common_chat_params common_chat_params_init_openpangu(const common_chat_template & tmpl, const autoparser::generation_params & inputs); ++ + common_chat_params common_chat_params_init_qwen3_coder(const common_chat_template & tmpl, const autoparser::generation_params & inputs); + + common_chat_params common_chat_params_init_translate_gemma(const common_chat_template & tmpl, const autoparser::generation_params & inputs); +diff --git a/common/parsers/sources.cmake b/common/parsers/sources.cmake +index 152397074..ca4b11654 100644 +--- a/common/parsers/sources.cmake ++++ b/common/parsers/sources.cmake +@@ -19,6 +19,7 @@ set(LLAMA_CHAT_PARSERS_SOURCES + ${CMAKE_CURRENT_LIST_DIR}/minimax-m3.cpp + ${CMAKE_CURRENT_LIST_DIR}/ministral3.cpp + ${CMAKE_CURRENT_LIST_DIR}/muse-glimmer.cpp ++ ${CMAKE_CURRENT_LIST_DIR}/openpangu.cpp + ${CMAKE_CURRENT_LIST_DIR}/qwen3-coder.cpp + ${CMAKE_CURRENT_LIST_DIR}/translate-gemma.cpp + ) +diff --git a/conversion/__init__.py b/conversion/__init__.py +index 8ce4cb733..4302cc1a4 100644 +--- a/conversion/__init__.py ++++ b/conversion/__init__.py +@@ -228,6 +228,7 @@ TEXT_MODEL_MAP: dict[str, str] = { + "PLaMo3ForCausalLM": "plamo", + "PaddleOCRVLForConditionalGeneration": "ernie", + "PanguEmbeddedForCausalLM": "pangu", ++ "OpenPanguV2ForCausalLM": "pangu", + "Phi3ForCausalLM": "phi", + "Phi4ForCausalLMV": "phi", + "PhiForCausalLM": "phi", +diff --git a/conversion/base.py b/conversion/base.py +index a5a80f27d..c5b389fb8 100644 +--- a/conversion/base.py ++++ b/conversion/base.py +@@ -2010,6 +2010,9 @@ class TextModel(ModelBase): + if chkhsh == "877081d19cf6996e2c4ff0e1236341e9b7bde288f5311a56a937f0afbbb3aeb5": + # ref: https://huggingface.co/deepseek-ai/DeepSeek-V3 + res = "deepseek-v3" ++ if chkhsh == "65df2fe396b537a53433301848c0a739f56d56f67ad3d35eba27961ac33c12bb": ++ # ref: https://huggingface.co/openpangu/openPangu-2.0-Flash ++ res = "openpangu-v2" + if chkhsh == "b3f499bb4255f8ca19fccd664443283318f2fd2414d5e0b040fbdd0cc195d6c5": + # ref: https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B + res = "deepseek-r1-qwen" +diff --git a/conversion/pangu.py b/conversion/pangu.py +index 74c76532b..ec2301459 100644 +--- a/conversion/pangu.py ++++ b/conversion/pangu.py +@@ -4,6 +4,8 @@ import json + + from typing import Iterable, TYPE_CHECKING + ++import torch ++ + if TYPE_CHECKING: + from torch import Tensor + +@@ -45,3 +47,228 @@ class PanguEmbeddedModel(TextModel): + logger.info("Skipping tied output layer 'lm_head.weight'") + return + yield from super().modify_tensors(data_torch, name, bid) ++ ++ ++@ModelBase.register("OpenPanguV2ForCausalLM") ++class OpenPanguV2Model(TextModel): ++ model_arch = gguf.MODEL_ARCH.OPENPANGU_V2 ++ ++ # simple root tensor map (HF name -> gguf enum) ++ _root_map = { ++ "model.embed_tokens.weight": (gguf.MODEL_TENSOR.TOKEN_EMBD, ".weight"), ++ "model.norm.weight": (gguf.MODEL_TENSOR.OUTPUT_NORM, ".weight"), ++ "lm_head.weight": (gguf.MODEL_TENSOR.OUTPUT, ".weight"), ++ "model.merge_mhc_module.branch_alpha_pre": (gguf.MODEL_TENSOR.HC_HEAD_SCALE, ".weight"), ++ "model.merge_mhc_module.branch_beta_pre": (gguf.MODEL_TENSOR.HC_HEAD_BASE, ".weight"), ++ } ++ ++ # per-layer tensor map (suffix after "model.layers.N." -> gguf enum) ++ _layer_map = { ++ "input_layernorm.weight": (gguf.MODEL_TENSOR.ATTN_NORM, ".weight"), ++ "post_attention_layernorm.weight": (gguf.MODEL_TENSOR.ATTN_POST_NORM, ".weight"), ++ "pre_mlp_layernorm.weight": (gguf.MODEL_TENSOR.FFN_NORM, ".weight"), ++ "post_mlp_layernorm.weight": (gguf.MODEL_TENSOR.FFN_POST_NORM, ".weight"), ++ "block_post_layernorm.weight": (gguf.MODEL_TENSOR.BLOCK_POST_NORM, ".weight"), ++ "attn_mhc_module.branch_alpha": (gguf.MODEL_TENSOR.HC_ATTN_SCALE, ".weight"), ++ "attn_mhc_module.branch_beta": (gguf.MODEL_TENSOR.HC_ATTN_BASE, ".weight"), ++ "mlp_mhc_module.branch_alpha": (gguf.MODEL_TENSOR.HC_FFN_SCALE, ".weight"), ++ "mlp_mhc_module.branch_beta": (gguf.MODEL_TENSOR.HC_FFN_BASE, ".weight"), ++ "self_attn.q_a_proj.weight": (gguf.MODEL_TENSOR.ATTN_Q_A, ".weight"), ++ "self_attn.q_a_layernorm.weight": (gguf.MODEL_TENSOR.ATTN_Q_A_NORM, ".weight"), ++ "self_attn.q_b_proj.weight": (gguf.MODEL_TENSOR.ATTN_Q_B, ".weight"), ++ "self_attn.qa_conv.weight": (gguf.MODEL_TENSOR.ATTN_Q_A_CONV, ".weight"), ++ "self_attn.kv_a_proj_with_mqa.weight": (gguf.MODEL_TENSOR.ATTN_KV_A_MQA, ".weight"), ++ "self_attn.kv_a_layernorm.weight": (gguf.MODEL_TENSOR.ATTN_KV_A_NORM, ".weight"), ++ "self_attn.compresskv_conv.weight": (gguf.MODEL_TENSOR.ATTN_KV_A_CONV, ".weight"), ++ "self_attn.o_conv.weight": (gguf.MODEL_TENSOR.ATTN_O_CONV, ".weight"), ++ "self_attn.o_proj.weight": (gguf.MODEL_TENSOR.ATTN_OUT, ".weight"), ++ "self_attn.param_sink_compressed_kv": (gguf.MODEL_TENSOR.ATTN_SINK_KV, ".weight"), ++ "self_attn.param_sink_k_pe": (gguf.MODEL_TENSOR.ATTN_SINK_K_PE, ".weight"), ++ "self_attn.indexer.wq_b.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_Q_B, ".weight"), ++ "self_attn.indexer.wk.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_K, ".weight"), ++ "self_attn.indexer.k_norm.weight": (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".weight"), ++ "self_attn.indexer.weights_proj.weight": (gguf.MODEL_TENSOR.INDEXER_PROJ, ".weight"), ++ "mlp.gate.weight": (gguf.MODEL_TENSOR.FFN_GATE_INP, ".weight"), ++ "mlp.e_score_correction.bias": (gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, ".bias"), ++ "mlp.shared_experts.gate_proj.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"), ++ "mlp.shared_experts.up_proj.weight": (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"), ++ "mlp.shared_experts.down_proj.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"), ++ "mlp.gate_proj.weight": (gguf.MODEL_TENSOR.FFN_GATE, ".weight"), ++ "mlp.up_proj.weight": (gguf.MODEL_TENSOR.FFN_UP, ".weight"), ++ "mlp.down_proj.weight": (gguf.MODEL_TENSOR.FFN_DOWN, ".weight"), ++ # MTP / NextN heads (layers >= num_hidden_layers) ++ "enorm.weight": (gguf.MODEL_TENSOR.NEXTN_ENORM, ".weight"), ++ "hnorm.weight": (gguf.MODEL_TENSOR.NEXTN_HNORM, ".weight"), ++ "eh_proj.weight": (gguf.MODEL_TENSOR.NEXTN_EH_PROJ, ".weight"), ++ "embed_tokens.weight": (gguf.MODEL_TENSOR.NEXTN_EMBED_TOKENS, ".weight"), ++ "shared_head.norm.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ".weight"), ++ "shared_head.head.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, ".weight"), ++ } ++ ++ _expert_map = { ++ "gate_proj": gguf.MODEL_TENSOR.FFN_GATE_EXP, ++ "up_proj": gguf.MODEL_TENSOR.FFN_UP_EXP, ++ "down_proj": gguf.MODEL_TENSOR.FFN_DOWN_EXP, ++ } ++ ++ def __init__(self, *args, **kwargs): ++ super().__init__(*args, **kwargs) ++ # buffers for mHC gamma-fold and expert stacking ++ self._mhc_buf: dict = {} ++ self._exp_buf: dict = {} ++ # include the MTP / NextN prediction layers in the block count ++ self._n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0) or 0) ++ if self._n_nextn > 0: ++ self.block_count = self.hparams["num_hidden_layers"] + self._n_nextn ++ self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) ++ ++ def tensor_force_quant(self, name, new_name, bid, n_dims): ++ # tiny depthwise conv kernels, consumed as F32 by ggml_ssm_conv ++ if new_name.endswith("_conv.weight") and ".attn_" in new_name: ++ return gguf.GGMLQuantizationType.F32 ++ return super().tensor_force_quant(name, new_name, bid, n_dims) ++ ++ def set_vocab(self): ++ # openPangu ships a custom tokenizer_class in tokenizer_config.json, but ++ # tokenizer.json is a standard ByteLevel BPE. Allow AutoTokenizer to load ++ # it without prompting for the custom-code confirmation. ++ from transformers import AutoTokenizer ++ orig = AutoTokenizer.from_pretrained ++ ++ def patched(*args, **kwargs): ++ kwargs.setdefault("trust_remote_code", True) ++ return orig(*args, **kwargs) ++ ++ AutoTokenizer.from_pretrained = staticmethod(patched) ++ try: ++ self._set_vocab_gpt2() ++ finally: ++ AutoTokenizer.from_pretrained = orig ++ ++ def set_gguf_parameters(self): ++ hparams = self.hparams ++ # MLA uses a single compressed KV head ++ hparams.setdefault("num_key_value_heads", 1) ++ ++ super().set_gguf_parameters() ++ ++ self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"]) ++ ++ self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"]) ++ self.gguf_writer.add_kv_lora_rank(hparams["kv_lora_rank"]) ++ self.gguf_writer.add_key_length_mla(hparams["qk_nope_head_dim"] + hparams["qk_rope_head_dim"]) ++ self.gguf_writer.add_value_length_mla(hparams["v_head_dim"]) ++ ++ self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"]) ++ self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) ++ self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"]) ++ self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"]) ++ self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"]) ++ self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) ++ ++ self.gguf_writer.add_indexer_head_count(hparams["index_n_heads"]) ++ self.gguf_writer.add_indexer_key_length(hparams["index_head_dim"]) ++ self.gguf_writer.add_indexer_top_k(hparams["index_topk"]) ++ ++ self.gguf_writer.add_hyper_connection_count(hparams["mhc_num_stream"]) ++ self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hparams["mhc_recur_norm"]) ++ self.gguf_writer.add_hyper_connection_epsilon(1e-6) ++ ++ self.gguf_writer.add_sink_count(hparams["param_sink_number"]) ++ self.gguf_writer.add_conv_kernel_size(hparams["router_sliding_window"]) ++ ++ if self._n_nextn > 0: ++ self.gguf_writer.add_nextn_predict_layers(self._n_nextn) ++ ++ # per-layer sliding window: DSA layers get 0, SWA layers get their window. ++ # NOTE: the reference runs the MTP / NextN layers with SWA (window 2048, ++ # sliding_window_list[-1]), but llama.cpp has a single global window size and ++ # the trunk SWA layers use 512, so the MTP layers are converted as global (0). ++ # This diverges from the reference at contexts beyond the MTP window and can ++ # only lower draft acceptance (drafts are verified by the trunk). ++ n_layer = hparams["num_hidden_layers"] ++ swa_layers = hparams["swa_layers"] ++ window_list = hparams["sliding_window_list"] ++ layer_window = [] ++ for il in range(self.block_count): ++ if il < n_layer and il in swa_layers: ++ layer_window.append(int(window_list[swa_layers.index(il)])) ++ else: ++ layer_window.append(0) ++ self.gguf_writer.add_sliding_window(int(hparams["sliding_window"])) ++ self.gguf_writer.add_layer_swa_window(layer_window) ++ ++ def _fold_mhc(self, key, tensor_enum, bid, data_torch, is_phi): ++ buf = self._mhc_buf.setdefault(key, {}) ++ buf["phi" if is_phi else "gamma"] = data_torch ++ if "phi" not in buf or "gamma" not in buf: ++ return [] ++ # fold RMSNorm gamma into phi: linear(x*gamma, W) == linear(x, W*gamma) ++ folded = (buf["phi"].float() * buf["gamma"].float()) ++ self._mhc_buf.pop(key) ++ return [(self.format_tensor_name(tensor_enum, bid), folded)] ++ ++ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: ++ # MLA absorption: split kv_b_proj into k_b (transposed) and v_b ++ if bid is not None and name.endswith("self_attn.kv_b_proj.weight"): ++ n_head = self.hparams["num_attention_heads"] ++ v_head_dim = self.hparams["v_head_dim"] ++ qk_nope = self.hparams["qk_nope_head_dim"] ++ kv_b = data_torch.view(n_head, v_head_dim + qk_nope, data_torch.shape[-1]) ++ k_b, v_b = torch.split(kv_b, [qk_nope, v_head_dim], dim=1) ++ k_b = k_b.transpose(1, 2).contiguous() ++ v_b = v_b.contiguous() ++ return [ ++ (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid), k_b), ++ (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid), v_b), ++ ] ++ ++ # mHC phi + norm_gamma get folded together ++ if name == "model.merge_mhc_module.phi.weight": ++ return self._fold_mhc("merge", gguf.MODEL_TENSOR.HC_HEAD_FN, None, data_torch, True) ++ if name == "model.merge_mhc_module.norm_gamma": ++ return self._fold_mhc("merge", gguf.MODEL_TENSOR.HC_HEAD_FN, None, data_torch, False) ++ if name in self._root_map: ++ enum, suffix = self._root_map[name] ++ return [(self.format_tensor_name(enum, None, suffix), data_torch)] ++ ++ if bid is not None: ++ suffix = name.split(f"model.layers.{bid}.", 1)[-1] ++ ++ for mod, fn_enum in (("attn_mhc_module", gguf.MODEL_TENSOR.HC_ATTN_FN), ++ ("mlp_mhc_module", gguf.MODEL_TENSOR.HC_FFN_FN)): ++ if suffix == f"{mod}.phi.weight": ++ return self._fold_mhc((bid, mod), fn_enum, bid, data_torch, True) ++ if suffix == f"{mod}.norm_gamma": ++ return self._fold_mhc((bid, mod), fn_enum, bid, data_torch, False) ++ ++ # routed experts: stack across the expert dimension ++ if ".mlp.experts." in name: ++ parts = suffix.split(".") # mlp experts E proj weight ++ eid = int(parts[2]) ++ proj = parts[3] ++ ebuf = self._exp_buf.setdefault((bid, proj), {}) ++ ebuf[eid] = data_torch ++ n_experts = self.hparams["n_routed_experts"] ++ if len(ebuf) < n_experts: ++ return [] ++ stacked = torch.stack([ebuf[i] for i in range(n_experts)], dim=0) ++ self._exp_buf.pop((bid, proj)) ++ return [(self.format_tensor_name(self._expert_map[proj], bid), stacked)] ++ ++ if suffix in self._layer_map: ++ enum, sfx = self._layer_map[suffix] ++ out = data_torch ++ # depthwise conv weights are [C, 1, k]; drop the singleton group dim. ++ # store as F32 so ggml_ssm_conv uses them without a per-eval cast ++ # (like Mamba's ssm_conv1d; they are tiny) ++ if enum in (gguf.MODEL_TENSOR.ATTN_Q_A_CONV, ++ gguf.MODEL_TENSOR.ATTN_KV_A_CONV, ++ gguf.MODEL_TENSOR.ATTN_O_CONV): ++ if out.ndim == 3: ++ out = out.squeeze(1) ++ out = out.float() ++ return [(self.format_tensor_name(enum, bid, sfx), out)] ++ ++ logger.warning(f"openpangu-v2: unmapped tensor '{name}'") ++ return [] +diff --git a/conversion/split_pangu_mtp.py b/conversion/split_pangu_mtp.py +new file mode 100644 +index 000000000..bd96af26d +--- /dev/null ++++ b/conversion/split_pangu_mtp.py +@@ -0,0 +1,166 @@ ++#!/usr/bin/env python3 ++"""Split a combined openPangu-2.0 (openpangu-v2) GGUF into: ++ ++ - a base GGUF without the MTP / NextN layers (smaller; use standalone), and ++ - an MTP-only GGUF holding just the NextN head layers plus the shared ++ embedding / output tensors, usable for self-speculative decoding via ++ `--model-draft --mtp`. ++ ++The MTP GGUF renumbers the NextN layers to blk.0.. and carries ++block_count == nextn_predict_layers, so hparams.n_layer() == 0 and every ++layer is an MTP layer. Its per-layer SWA window is set to the reference ++value (sliding_window_list[-1], 2048 for openPangu-2.0-Flash) -- the split ++model runs the MTP head with its trained sliding window, which the combined ++GGUF cannot express next to the trunk's window. ++ ++Example: ++ python conversion/split_pangu_mtp.py openPangu-2.0-Flash-BF16.gguf \ ++ --base openPangu-2.0-Flash-base-BF16.gguf \ ++ --mtp openPangu-2.0-Flash-mtp-BF16.gguf ++""" ++ ++from __future__ import annotations ++ ++import argparse ++import logging ++import re ++import sys ++from pathlib import Path ++ ++if 'NO_LOCAL_GGUF' not in __import__('os').environ: ++ sys.path.insert(1, str(Path(__file__).parent.parent / 'gguf-py')) ++ ++import gguf ++from gguf import GGUFReader, GGUFWriter, GGUFValueType # noqa: E402 ++ ++logger = logging.getLogger("split-pangu-mtp") ++ ++# non-blk tensors the MTP head graph reads (shared head/embedding fallbacks) ++MTP_SHARED_TENSORS = {"token_embd.weight", "output.weight", "output_norm.weight"} ++ ++BLK_RE = re.compile(r"^blk\.(\d+)\.(.+)$") ++ ++ ++def copy_kv(reader: GGUFReader, writer: GGUFWriter, overrides: dict, remove: set) -> None: ++ for field in reader.fields.values(): ++ if field.name == gguf.Keys.General.ARCHITECTURE or field.name.startswith('GGUF.'): ++ continue ++ if field.name in remove: ++ logger.debug('removing %s', field.name) ++ continue ++ ++ val_type = field.types[0] ++ sub_type = field.types[-1] if val_type == GGUFValueType.ARRAY else None ++ ++ value = field.contents() ++ if field.name in overrides: ++ value = overrides.pop(field.name) ++ logger.info('overriding %s -> %s', field.name, str(value)[:80]) ++ ++ writer.add_key_value(field.name, value, val_type, sub_type=sub_type) ++ ++ for key, (value, val_type) in overrides.items(): ++ logger.info('adding %s = %s', key, str(value)[:80]) ++ writer.add_key_value(key, value, val_type) ++ ++ ++def write_split(reader: GGUFReader, path: Path, arch: str, overrides: dict, remove: set, ++ keep_tensor) -> None: ++ writer = GGUFWriter(path, arch) ++ copy_kv(reader, writer, dict(overrides), remove) ++ ++ tensors = [] ++ for tensor in reader.tensors: ++ new_name = keep_tensor(tensor.name) ++ if new_name is None: ++ continue ++ tensors.append((new_name, tensor)) ++ writer.add_tensor_info(new_name, tensor.data.shape, tensor.data.dtype, tensor.data.nbytes, ++ tensor.tensor_type) ++ ++ writer.write_header_to_file() ++ writer.write_kv_data_to_file() ++ writer.write_ti_data_to_file() ++ ++ for i, (new_name, tensor) in enumerate(tensors): ++ logger.info('[%d/%d] writing %s (%s)', i + 1, len(tensors), new_name, path.name) ++ writer.write_tensor_data(tensor.data, tensor_endianess=reader.endianess) ++ ++ writer.close() ++ ++ ++def main() -> None: ++ parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ++ parser.add_argument("input", type=Path, help="combined openpangu-v2 GGUF") ++ parser.add_argument("--base", type=Path, help="output path for the base (no-MTP) GGUF") ++ parser.add_argument("--mtp", type=Path, help="output path for the MTP-only GGUF") ++ parser.add_argument("--mtp-swa-window", type=int, default=2048, ++ help="SWA window for the MTP layers (reference: sliding_window_list[-1], default 2048)") ++ parser.add_argument("--verbose", action="store_true") ++ args = parser.parse_args() ++ ++ logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) ++ ++ if not args.base and not args.mtp: ++ parser.error("nothing to do: pass --base and/or --mtp") ++ ++ reader = GGUFReader(args.input) ++ ++ arch = reader.get_field(gguf.Keys.General.ARCHITECTURE).contents() ++ if arch != "openpangu-v2": ++ parser.error(f"expected an openpangu-v2 GGUF, got architecture '{arch}'") ++ ++ n_layer_all = reader.get_field(f"{arch}.block_count").contents() ++ field_nextn = reader.get_field(f"{arch}.nextn_predict_layers") ++ n_nextn = field_nextn.contents() if field_nextn else 0 ++ if n_nextn <= 0: ++ parser.error("this GGUF has no MTP / NextN layers (nextn_predict_layers missing or 0)") ++ n_base = n_layer_all - n_nextn ++ ++ swa_field = reader.get_field(f"{arch}.attention.layer_swa_window") ++ swa_all = list(swa_field.contents()) if swa_field else [0] * n_layer_all ++ ++ logger.info("input: %d layers = %d base + %d MTP", n_layer_all, n_base, n_nextn) ++ ++ if args.base: ++ def keep_base(name: str): ++ m = BLK_RE.match(name) ++ if m and int(m.group(1)) >= n_base: ++ return None ++ return name ++ ++ write_split(reader, args.base, arch, ++ overrides={ ++ f"{arch}.block_count": n_base, ++ f"{arch}.attention.layer_swa_window": swa_all[:n_base], ++ }, ++ remove={f"{arch}.nextn_predict_layers"}, ++ keep_tensor=keep_base) ++ logger.info("base GGUF written to %s", args.base) ++ ++ if args.mtp: ++ def keep_mtp(name: str): ++ m = BLK_RE.match(name) ++ if m: ++ il = int(m.group(1)) ++ if il < n_base: ++ return None ++ return f"blk.{il - n_base}.{m.group(2)}" ++ if name in MTP_SHARED_TENSORS: ++ return name ++ return None ++ ++ write_split(reader, args.mtp, arch, ++ overrides={ ++ f"{arch}.block_count": n_nextn, ++ f"{arch}.nextn_predict_layers": n_nextn, ++ f"{arch}.leading_dense_block_count": 0, ++ f"{arch}.attention.layer_swa_window": [args.mtp_swa_window] * n_nextn, ++ }, ++ remove=set(), ++ keep_tensor=keep_mtp) ++ logger.info("MTP GGUF written to %s", args.mtp) ++ ++ ++if __name__ == '__main__': ++ main() +diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h +index c37ba9471..efff904ab 100644 +--- a/ggml/include/ggml.h ++++ b/ggml/include/ggml.h +@@ -567,6 +567,9 @@ extern "C" { + GGML_OP_LEAKY_RELU, + GGML_OP_TRI, + GGML_OP_FILL, ++ GGML_OP_SINKHORN, ++ GGML_OP_HC_MIX, ++ GGML_OP_DSA_SCORE, + + GGML_OP_FLASH_ATTN_EXT, + GGML_OP_FLASH_ATTN_BACK, +@@ -2439,6 +2442,47 @@ extern "C" { + struct ggml_tensor * a, + float c); + ++ // Sinkhorn row/column normalization over the leading [ne0, ne1] matrix of each ++ // (i2, i3) batch (used by hyper-connection stream mixing): ++ // 1. softmax along dim0 (per i1) ++ // 2. add eps to every element ++ // 3. normalize along dim1: M[i0,:] /= (eps + sum_i1 M[i0,i1]) ++ // 4. repeat (n_iter-1) times: ++ // normalize along dim0: M[:,i1] /= (eps + sum_i0 M[i0,i1]), then step 3 ++ GGML_API struct ggml_tensor * ggml_sinkhorn( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ int n_iter, ++ float eps); ++ ++ // Hyper-connection stream mixing: for each token (rows of a [2*hc+hc*hc, nt]): ++ // out[0 .. hc) = sigmoid(a*scale[0] + base[i]) + eps (pre) ++ // out[hc .. 2*hc) = 2*sigmoid(a*scale[1] + base[i]) (post) ++ // out[2*hc .. 2*hc+hc*hc) = sinkhorn(a*scale[2] + base[i], n_iter, eps) (comb, ++ // the [hc, hc] matrix laid out dim0-major, same semantics as ggml_sinkhorn) ++ GGML_API struct ggml_tensor * ggml_hc_mix( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * scale, // [3] ++ struct ggml_tensor * base, // [2*hc+hc*hc] ++ int hc, ++ int n_iter, ++ float eps); ++ ++ // DSA lightning-indexer scores (one pass over the cached indexer keys): ++ // out[s,t] = mask[s,t] + sum_h w[h,t] * relu(sum_d ik[d,s] * q[d,h,t]) ++ // ik: [d, n_kv] indexer keys (may be a strided view; F16 or F32) ++ // q: [d, h, nt] per-token indexer queries (F32) ++ // w: [h, nt] per-token head weights (F32) ++ // mask: [n_kv, nt] additive mask (F16 or F32) ++ // returns [n_kv, nt] F32 ++ GGML_API struct ggml_tensor * ggml_dsa_score( ++ struct ggml_context * ctx, ++ struct ggml_tensor * ik, ++ struct ggml_tensor * q, ++ struct ggml_tensor * w, ++ struct ggml_tensor * mask); ++ + // Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151 + // timesteps: [N,] + // return: [N, dim] +diff --git a/ggml/src/ggml-cpu/amx/amx.cpp b/ggml/src/ggml-cpu/amx/amx.cpp +index e90c2797f..37780d618 100644 +--- a/ggml/src/ggml-cpu/amx/amx.cpp ++++ b/ggml/src/ggml-cpu/amx/amx.cpp +@@ -146,6 +146,9 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type { + if (op->op != GGML_OP_MUL_MAT) { + return false; + } ++ if (op->src[0]->ne[2] > 1 || op->src[0]->ne[3] > 1) { ++ return false; // packing is not implemented for batched tensors ++ } + auto * src0 = op->src[0]; + auto * src1 = op->src[1]; + +diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c +index ab82e3d6c..5d13dbea1 100644 +--- a/ggml/src/ggml-cpu/ggml-cpu.c ++++ b/ggml/src/ggml-cpu/ggml-cpu.c +@@ -2026,6 +2026,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm + { + ggml_compute_forward_fill(params, tensor); + } break; ++ case GGML_OP_SINKHORN: ++ { ++ ggml_compute_forward_sinkhorn(params, tensor); ++ } break; ++ case GGML_OP_HC_MIX: ++ { ++ ggml_compute_forward_hc_mix(params, tensor); ++ } break; ++ case GGML_OP_DSA_SCORE: ++ { ++ ggml_compute_forward_dsa_score(params, tensor); ++ } break; + case GGML_OP_FLASH_ATTN_EXT: + { + ggml_compute_forward_flash_attn_ext(params, tensor); +@@ -2266,6 +2278,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { + case GGML_OP_CUMSUM: + case GGML_OP_TRI: + case GGML_OP_FILL: ++ case GGML_OP_SINKHORN: ++ case GGML_OP_HC_MIX: ++ case GGML_OP_DSA_SCORE: + { + n_tasks = n_threads; + } break; +diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp +index 7a73e138e..56c864e39 100644 +--- a/ggml/src/ggml-cpu/ops.cpp ++++ b/ggml/src/ggml-cpu/ops.cpp +@@ -2334,6 +2334,202 @@ void ggml_compute_forward_fill(const ggml_compute_params * params, ggml_tensor * + } + } + ++// ggml_compute_forward_sinkhorn ++ ++// shared by GGML_OP_SINKHORN and GGML_OP_HC_MIX: softmax along dim0, add eps, ++// then alternating dim1/dim0 normalization of the [ne0, ne1] matrix m ++static void ggml_sinkhorn_matrix(float * m, const int64_t ne0, const int64_t ne1, const int n_iter, const float eps) { ++ const int64_t nm = ne0*ne1; ++ ++ // softmax along dim0, per i1 ++ for (int64_t i1 = 0; i1 < ne1; ++i1) { ++ float * c = m + i1*ne0; ++ float vmax = c[0]; ++ for (int64_t i0 = 1; i0 < ne0; ++i0) { ++ vmax = MAX(vmax, c[i0]); ++ } ++ float sum = 0.0f; ++ for (int64_t i0 = 0; i0 < ne0; ++i0) { ++ c[i0] = expf(c[i0] - vmax); ++ sum += c[i0]; ++ } ++ // reciprocal multiply to match ggml_soft_max rounding ++ const float inv = 1.0f/sum; ++ for (int64_t i0 = 0; i0 < ne0; ++i0) { ++ c[i0] *= inv; ++ } ++ } ++ ++ // add eps to every element ++ for (int64_t i = 0; i < nm; ++i) { ++ m[i] += eps; ++ } ++ ++ for (int it = 0; it < n_iter; ++it) { ++ // normalize along dim0, per i1 (skipped on the first pass) ++ if (it > 0) { ++ for (int64_t i1 = 0; i1 < ne1; ++i1) { ++ float sum = eps; ++ for (int64_t i0 = 0; i0 < ne0; ++i0) { ++ sum += m[i0 + i1*ne0]; ++ } ++ for (int64_t i0 = 0; i0 < ne0; ++i0) { ++ m[i0 + i1*ne0] /= sum; ++ } ++ } ++ } ++ // normalize along dim1, per i0 ++ for (int64_t i0 = 0; i0 < ne0; ++i0) { ++ float sum = eps; ++ for (int64_t i1 = 0; i1 < ne1; ++i1) { ++ sum += m[i0 + i1*ne0]; ++ } ++ for (int64_t i1 = 0; i1 < ne1; ++i1) { ++ m[i0 + i1*ne0] /= sum; ++ } ++ } ++ } ++} ++ ++void ggml_compute_forward_sinkhorn(const ggml_compute_params * params, ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; ++ ++ GGML_ASSERT(src0->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_is_contiguous(src0)); ++ GGML_ASSERT(ggml_is_contiguous(dst)); ++ ++ const int n_iter = ggml_get_op_params_i32(dst, 0); ++ const float eps = ggml_get_op_params_f32(dst, 1); ++ ++ const int64_t ne0 = src0->ne[0]; ++ const int64_t ne1 = src0->ne[1]; ++ const int64_t nm = ne0*ne1; // per-batch matrix size (<= 64, asserted at build) ++ const int64_t nb = src0->ne[2]*src0->ne[3]; ++ ++ // parallelize over the (i2, i3) batches ++ const int64_t db = (nb + params->nth - 1)/params->nth; ++ const int64_t b0 = db*params->ith; ++ const int64_t b1 = MIN(b0 + db, nb); ++ ++ float m[64]; ++ ++ for (int64_t b = b0; b < b1; ++b) { ++ const float * sp = (const float *) src0->data + b*nm; ++ float * dp = (float *) dst->data + b*nm; ++ ++ memcpy(m, sp, nm*sizeof(float)); ++ ++ ggml_sinkhorn_matrix(m, ne0, ne1, n_iter, eps); ++ ++ memcpy(dp, m, nm*sizeof(float)); ++ } ++} ++ ++// ggml_compute_forward_hc_mix ++ ++void ggml_compute_forward_hc_mix(const ggml_compute_params * params, ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; // mixes [2*hc + hc*hc, nt] ++ const ggml_tensor * scale = dst->src[1]; // [3] ++ const ggml_tensor * base = dst->src[2]; // [2*hc + hc*hc] ++ ++ GGML_ASSERT(src0->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_is_contiguous(src0)); ++ GGML_ASSERT(ggml_is_contiguous(dst)); ++ ++ const int hc = ggml_get_op_params_i32(dst, 0); ++ const int n_iter = ggml_get_op_params_i32(dst, 1); ++ const float eps = ggml_get_op_params_f32(dst, 2); ++ ++ const int64_t n0 = src0->ne[0]; ++ const int64_t nt = src0->ne[1]*src0->ne[2]*src0->ne[3]; ++ ++ GGML_ASSERT(n0 == 2*hc + hc*hc); ++ ++ const float * sc = (const float *) scale->data; ++ const float * bs = (const float *) base->data; ++ ++ // parallelize over the tokens ++ const int64_t dt = (nt + params->nth - 1)/params->nth; ++ const int64_t t0 = dt*params->ith; ++ const int64_t t1 = MIN(t0 + dt, nt); ++ ++ float m[64]; ++ ++ for (int64_t t = t0; t < t1; ++t) { ++ const float * sp = (const float *) src0->data + t*n0; ++ float * dp = (float *) dst->data + t*n0; ++ ++ for (int i = 0; i < hc; ++i) { ++ dp[i] = 1.0f/(1.0f + expf(-(sp[i]*sc[0] + bs[i]))) + eps; // pre ++ } ++ for (int i = hc; i < 2*hc; ++i) { ++ dp[i] = 2.0f/(1.0f + expf(-(sp[i]*sc[1] + bs[i]))); // post ++ } ++ for (int i = 2*hc; i < (int) n0; ++i) { ++ m[i - 2*hc] = sp[i]*sc[2] + bs[i]; // comb ++ } ++ ++ ggml_sinkhorn_matrix(m, hc, hc, n_iter, eps); ++ ++ memcpy(dp + 2*hc, m, hc*hc*sizeof(float)); ++ } ++} ++ ++// ggml_compute_forward_dsa_score ++ ++void ggml_compute_forward_dsa_score(const ggml_compute_params * params, ggml_tensor * dst) { ++ const ggml_tensor * ik = dst->src[0]; // [d, n_kv] (possibly strided rows) ++ const ggml_tensor * q = dst->src[1]; // [d, h, nt] ++ const ggml_tensor * w = dst->src[2]; // [h, nt] ++ const ggml_tensor * mask = dst->src[3]; // [n_kv, nt] ++ ++ const int64_t nd = ik->ne[0]; ++ const int64_t nkv = ik->ne[1]; ++ const int64_t nh = q->ne[1]; ++ const int64_t nt = q->ne[2]; ++ ++ const float * qd = (const float *) q->data; ++ const float * wd = (const float *) w->data; ++ ++ // parallelize over the cached positions ++ const int64_t ds = (nkv + params->nth - 1)/params->nth; ++ const int64_t s0 = ds*params->ith; ++ const int64_t s1 = MIN(s0 + ds, nkv); ++ ++ for (int64_t s = s0; s < s1; ++s) { ++ const char * ik_row = (const char *) ik->data + s*ik->nb[1]; ++ ++ for (int64_t t = 0; t < nt; ++t) { ++ float acc = 0.0f; ++ for (int64_t h = 0; h < nh; ++h) { ++ const float * qh = qd + (t*nh + h)*nd; ++ float dot = 0.0f; ++ if (ik->type == GGML_TYPE_F16) { ++ const ggml_fp16_t * ikd = (const ggml_fp16_t *) ik_row; ++ for (int64_t d = 0; d < nd; ++d) { ++ dot += GGML_CPU_FP16_TO_FP32(ikd[d])*qh[d]; ++ } ++ } else { ++ const float * ikd = (const float *) ik_row; ++ for (int64_t d = 0; d < nd; ++d) { ++ dot += ikd[d]*qh[d]; ++ } ++ } ++ acc += wd[t*nh + h]*(dot > 0.0f ? dot : 0.0f); ++ } ++ ++ float m; ++ if (mask->type == GGML_TYPE_F16) { ++ m = GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) ((const char *) mask->data + t*mask->nb[1]))[s]); ++ } else { ++ m = ((const float *) ((const char *) mask->data + t*mask->nb[1]))[s]; ++ } ++ ++ ((float *) dst->data)[t*nkv + s] = acc + m; ++ } ++ } ++} ++ + // ggml_compute_tri + + static void ggml_compute_forward_tri_f32(const ggml_compute_params * params, ggml_tensor * dst) { +diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h +index 2728b08b6..58821c0c7 100644 +--- a/ggml/src/ggml-cpu/ops.h ++++ b/ggml/src/ggml-cpu/ops.h +@@ -86,6 +86,9 @@ void ggml_compute_forward_top_k(const struct ggml_compute_params * params, struc + void ggml_compute_forward_leaky_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_sinkhorn(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_hc_mix(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_dsa_score(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_flash_attn_back( + const struct ggml_compute_params * params, +diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu +index f8dcd95ab..7286e76eb 100644 +--- a/ggml/src/ggml-cuda/concat.cu ++++ b/ggml/src/ggml-cuda/concat.cu +@@ -1,5 +1,6 @@ + #include "concat.cuh" + ++#include + #include + + // contiguous kernels +@@ -103,39 +104,45 @@ static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) + uint64_t nb12, + uint64_t nb13, + int64_t ne0, +- int64_t /*ne1*/, +- int64_t /*ne2*/, +- int64_t /*ne3*/, ++ int64_t ne1, ++ int64_t ne2, ++ int64_t ne3, + uint64_t nb0, + uint64_t nb1, + uint64_t nb2, + uint64_t nb3) { + static_assert(dim >= 0 && dim <= 3, "dim must be in [0, 3]"); + +- const int64_t i3 = blockIdx.z; +- const int64_t i2 = blockIdx.y; +- const int64_t i1 = blockIdx.x; ++ // rows are flattened into a grid-stride loop over gridDim.x because ++ // ne1/ne2/ne3 can each exceed the CUDA grid.y/grid.z limit of 65535 ++ const int64_t n_rows = ne1*ne2*ne3; + + const T * x; + +- for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) { +- if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { +- x = (const T *)(src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); +- } else { +- if constexpr (dim == 0) { +- x = (const T *)(src1 + i3*nb13 + i2*nb12 + i1*nb11 + (i0 - ne00)*nb10); +- } else if constexpr (dim == 1) { +- x = (const T *)(src1 + i3*nb13 + i2*nb12 + (i1 - ne01)*nb11 + i0*nb10); +- } else if constexpr (dim == 2) { +- x = (const T *)(src1 + i3*nb13 + (i2 - ne02)*nb12 + i1*nb11 + i0*nb10); +- } else if constexpr (dim == 3) { +- x = (const T *)(src1 + (i3 - ne03)*nb13 + i2*nb12 + i1*nb11 + i0*nb10); ++ for (int64_t ir = blockIdx.x; ir < n_rows; ir += gridDim.x) { ++ const int64_t i3 = ir / (ne1*ne2); ++ const int64_t i2 = (ir - i3*ne1*ne2) / ne1; ++ const int64_t i1 = ir - i3*ne1*ne2 - i2*ne1; ++ ++ for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) { ++ if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { ++ x = (const T *)(src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); ++ } else { ++ if constexpr (dim == 0) { ++ x = (const T *)(src1 + i3*nb13 + i2*nb12 + i1*nb11 + (i0 - ne00)*nb10); ++ } else if constexpr (dim == 1) { ++ x = (const T *)(src1 + i3*nb13 + i2*nb12 + (i1 - ne01)*nb11 + i0*nb10); ++ } else if constexpr (dim == 2) { ++ x = (const T *)(src1 + i3*nb13 + (i2 - ne02)*nb12 + i1*nb11 + i0*nb10); ++ } else if constexpr (dim == 3) { ++ x = (const T *)(src1 + (i3 - ne03)*nb13 + i2*nb12 + i1*nb11 + i0*nb10); ++ } + } +- } + +- T * y = (T *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); ++ T * y = (T *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); + +- *y = *x; ++ *y = *x; ++ } + } + } + +@@ -163,7 +170,8 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml + } else { + GGML_ASSERT(!ggml_is_quantized(src0->type)); + +- dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); ++ const int64_t n_rows = dst->ne[1]*dst->ne[2]*dst->ne[3]; ++ dim3 grid_dim((unsigned) std::min(n_rows, INT32_MAX)); + auto launch_kernel = [&](auto dim) { + concat_non_cont<<>>( + (const char *) src0->data, (const char *) src1->data, (char *) dst->data, +diff --git a/ggml/src/ggml-cuda/dsa-score.cu b/ggml/src/ggml-cuda/dsa-score.cu +new file mode 100644 +index 000000000..bf0a057f4 +--- /dev/null ++++ b/ggml/src/ggml-cuda/dsa-score.cu +@@ -0,0 +1,217 @@ ++#include "dsa-score.cuh" ++ ++// DSA lightning-indexer scores, one pass over the cached indexer keys: ++// out[s,t] = mask[s,t] + sum_h w[h,t] * relu(sum_d ik[d,s] * q[d,h,t]) ++// One warp per cached position s: the row load is coalesced across the lanes and ++// the per-head dot products are warp-shuffle reductions. q and w for the block's ++// token are staged in shared memory (d*h floats, <= 12 KB for the openPangu ++// 128x24 indexer). Used on the decode path (nt small); prefill keeps the tiled ++// matmul pipeline. ++template ++static __global__ void dsa_score_f32( ++ const char * ik, const float * q, const float * w, const char * mask, float * dst, ++ const int nd, const int64_t nkv, const int nh, const int64_t nb_ik_row, const int64_t nb_mask_row) { ++ extern __shared__ float smem[]; // [nd*nh + nh] ++ ++ const int t = blockIdx.y; ++ const int warp = threadIdx.x/WARP_SIZE; ++ const int lane = threadIdx.x%WARP_SIZE; ++ ++ const int warps_per_block = blockDim.x/WARP_SIZE; ++ const int64_t s = (int64_t) blockIdx.x*warps_per_block + warp; ++ ++ float * q_s = smem; ++ float * w_s = smem + (size_t) nd*nh; ++ ++ for (int i = threadIdx.x; i < nd*nh; i += blockDim.x) { ++ q_s[i] = q[(size_t) t*nd*nh + i]; ++ } ++ for (int i = threadIdx.x; i < nh; i += blockDim.x) { ++ w_s[i] = w[(size_t) t*nh + i]; ++ } ++ __syncthreads(); ++ ++ if (s >= nkv) { ++ return; ++ } ++ ++ const T_IK * ik_row = (const T_IK *) (ik + s*nb_ik_row); ++ ++ // coalesced row load: each lane keeps its d-strided slice in registers ++ float ikl[8]; // nd <= 8*WARP_SIZE (256), asserted host-side ++ const int per_lane = (nd + WARP_SIZE - 1)/WARP_SIZE; ++ for (int i = 0; i < per_lane; ++i) { ++ const int d = lane + i*WARP_SIZE; ++ ikl[i] = d < nd ? (float) ik_row[d] : 0.0f; ++ } ++ ++ float acc = 0.0f; ++ for (int h = 0; h < nh; ++h) { ++ const float * qh = q_s + h*nd; ++ float part = 0.0f; ++ for (int i = 0; i < per_lane; ++i) { ++ const int d = lane + i*WARP_SIZE; ++ if (d < nd) { ++ part += ikl[i]*qh[d]; ++ } ++ } ++ part = warp_reduce_sum(part); ++ acc += w_s[h]*fmaxf(part, 0.0f); ++ } ++ ++ if (lane == 0) { ++ const T_MASK * mask_row = (const T_MASK *) (mask + t*nb_mask_row); ++ dst[(size_t) t*nkv + s] = acc + (float) mask_row[s]; ++ } ++} ++ ++// Prefill variant: GEMM-like tiled kernel for large token counts. Each block ++// computes a DSA_TILE_S x DSA_TILE_T output tile; the (transposed) ik tile is ++// staged in shared memory once and reused across all heads, the per-head q tile ++// is restaged each head iteration. Each thread accumulates a 4x2 register tile, ++// so the shared-memory loads amortize over 8 FMAs. Never materializes the ++// [n_kv, nh, nt] intermediate the unfused matmul pipeline needs (~5 GB per ++// layer per 512-token ubatch at 100K context). ++#define DSA_TILE_S 64 ++#define DSA_TILE_T 32 ++#define DSA_TILE_PS (DSA_TILE_S + 4) // padded rows: keeps 8-byte alignment, avoids bank conflicts ++#define DSA_TILE_PT (DSA_TILE_T + 2) ++#define DSA_TILE_THREADS 256 ++#define DSA_TILE_ND_MAX 180 // smem = nd*(PS*2 + PT*4) bytes <= 48 KB ++ ++template ++static __global__ void dsa_score_tiled_f32( ++ const char * ik, const float * q, const float * w, const char * mask, float * dst, ++ const int nd, const int64_t nkv, const int nh, const int nt, ++ const int64_t nb_ik_row, const int64_t nb_mask_row) { ++ extern __shared__ char smem_raw[]; ++ half * sik = (half *) smem_raw; // [nd][DSA_TILE_PS], transposed ++ float * sq = (float *) (smem_raw + (size_t) nd*DSA_TILE_PS*sizeof(half)); // [nd][DSA_TILE_PT], transposed ++ ++ const int64_t s0 = (int64_t) blockIdx.x*DSA_TILE_S; ++ const int t0 = blockIdx.y*DSA_TILE_T; ++ ++ const int tid = threadIdx.x; ++ const int tt = tid % (DSA_TILE_T/2); // covers t = t0 + tt*2 + {0,1} ++ const int ts = tid / (DSA_TILE_T/2); // covers s = s0 + ts*4 + {0..3} ++ ++ // ik tile: coalesced global reads (d contiguous per row), transposed store ++ for (int idx = tid; idx < nd*DSA_TILE_S; idx += DSA_TILE_THREADS) { ++ const int d = idx % nd; ++ const int s = idx / nd; ++ sik[d*DSA_TILE_PS + s] = s0 + s < nkv ? ((const half *) (ik + (s0 + s)*nb_ik_row))[d] : __float2half(0.0f); ++ } ++ ++ float acc[4][2] = {{0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}}; ++ ++ for (int h = 0; h < nh; ++h) { ++ __syncthreads(); // also covers the sik stores on the first iteration ++ for (int idx = tid; idx < nd*DSA_TILE_T; idx += DSA_TILE_THREADS) { ++ const int d = idx % nd; ++ const int t = idx / nd; ++ sq[d*DSA_TILE_PT + t] = t0 + t < nt ? q[((size_t) (t0 + t)*nh + h)*nd + d] : 0.0f; ++ } ++ __syncthreads(); ++ ++ float dot[4][2] = {{0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}}; ++ for (int d = 0; d < nd; ++d) { ++ const half2 * a = (const half2 *) (sik + d*DSA_TILE_PS + ts*4); ++ const float2 a01 = __half22float2(a[0]); ++ const float2 a23 = __half22float2(a[1]); ++ const float2 b = *(const float2 *) (sq + d*DSA_TILE_PT + tt*2); ++ dot[0][0] += a01.x*b.x; dot[0][1] += a01.x*b.y; ++ dot[1][0] += a01.y*b.x; dot[1][1] += a01.y*b.y; ++ dot[2][0] += a23.x*b.x; dot[2][1] += a23.x*b.y; ++ dot[3][0] += a23.y*b.x; dot[3][1] += a23.y*b.y; ++ } ++ ++#pragma unroll ++ for (int j = 0; j < 2; ++j) { ++ const int t = t0 + tt*2 + j; ++ const float wv = t < nt ? w[(size_t) t*nh + h] : 0.0f; ++#pragma unroll ++ for (int i = 0; i < 4; ++i) { ++ acc[i][j] += wv*fmaxf(dot[i][j], 0.0f); ++ } ++ } ++ } ++ ++#pragma unroll ++ for (int j = 0; j < 2; ++j) { ++ const int t = t0 + tt*2 + j; ++ if (t >= nt) { ++ continue; ++ } ++ const T_MASK * mask_row = (const T_MASK *) (mask + (size_t) t*nb_mask_row); ++#pragma unroll ++ for (int i = 0; i < 4; ++i) { ++ const int64_t s = s0 + ts*4 + i; ++ if (s < nkv) { ++ dst[(size_t) t*nkv + s] = acc[i][j] + (float) mask_row[s]; ++ } ++ } ++ } ++} ++ ++void ggml_cuda_op_dsa_score(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * ik = dst->src[0]; ++ const ggml_tensor * q = dst->src[1]; ++ const ggml_tensor * w = dst->src[2]; ++ const ggml_tensor * mask = dst->src[3]; ++ ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ GGML_ASSERT(q->type == GGML_TYPE_F32 && w->type == GGML_TYPE_F32); ++ ++ const int nd = (int) ik->ne[0]; ++ const int64_t nkv = ik->ne[1]; ++ const int nh = (int) q->ne[1]; ++ const int nt = (int) q->ne[2]; ++ ++ GGML_ASSERT(nd <= 8*WARP_SIZE); ++ ++ // large token counts (prefill): tiled kernel, one output tile per block ++ if (nt >= 16 && ik->type == GGML_TYPE_F16 && nd <= DSA_TILE_ND_MAX) { ++ const dim3 grid((nkv + DSA_TILE_S - 1)/DSA_TILE_S, (nt + DSA_TILE_T - 1)/DSA_TILE_T, 1); ++ const size_t smem = (size_t) nd*(DSA_TILE_PS*sizeof(half) + DSA_TILE_PT*sizeof(float)); ++ ++ if (mask->type == GGML_TYPE_F16) { ++ dsa_score_tiled_f32<<>>( ++ (const char *) ik->data, (const float *) q->data, (const float *) w->data, ++ (const char *) mask->data, (float *) dst->data, ++ nd, nkv, nh, nt, ik->nb[1], mask->nb[1]); ++ } else { ++ dsa_score_tiled_f32<<>>( ++ (const char *) ik->data, (const float *) q->data, (const float *) w->data, ++ (const char *) mask->data, (float *) dst->data, ++ nd, nkv, nh, nt, ik->nb[1], mask->nb[1]); ++ } ++ return; ++ } ++ ++ const int warps_per_block = CUDA_DSA_SCORE_BLOCK_SIZE/WARP_SIZE; ++ const dim3 grid((nkv + warps_per_block - 1)/warps_per_block, nt, 1); ++ const size_t smem = ((size_t) nd*nh + nh)*sizeof(float); ++ ++ const auto launch = [&](auto ik_dummy, auto mask_dummy) { ++ using T_IK = decltype(ik_dummy); ++ using T_MASK = decltype(mask_dummy); ++ dsa_score_f32<<>>( ++ (const char *) ik->data, (const float *) q->data, (const float *) w->data, ++ (const char *) mask->data, (float *) dst->data, ++ nd, nkv, nh, ik->nb[1], mask->nb[1]); ++ }; ++ ++ if (ik->type == GGML_TYPE_F16) { ++ if (mask->type == GGML_TYPE_F16) { ++ launch(half{}, half{}); ++ } else { ++ launch(half{}, float{}); ++ } ++ } else { ++ if (mask->type == GGML_TYPE_F16) { ++ launch(float{}, half{}); ++ } else { ++ launch(float{}, float{}); ++ } ++ } ++} +diff --git a/ggml/src/ggml-cuda/dsa-score.cuh b/ggml/src/ggml-cuda/dsa-score.cuh +new file mode 100644 +index 000000000..564a4ea92 +--- /dev/null ++++ b/ggml/src/ggml-cuda/dsa-score.cuh +@@ -0,0 +1,5 @@ ++#include "common.cuh" ++ ++#define CUDA_DSA_SCORE_BLOCK_SIZE 128 ++ ++void ggml_cuda_op_dsa_score(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu +index 7c45d0d1f..b835759f7 100644 +--- a/ggml/src/ggml-cuda/ggml-cuda.cu ++++ b/ggml/src/ggml-cuda/ggml-cuda.cu +@@ -70,6 +70,8 @@ + #include "ggml-cuda/cumsum.cuh" + #include "ggml-cuda/fill.cuh" + #include "ggml-cuda/lightning-indexer.cuh" ++#include "ggml-cuda/sinkhorn.cuh" ++#include "ggml-cuda/dsa-score.cuh" + #include "ggml.h" + + #include +@@ -2542,6 +2544,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg + case GGML_OP_LIGHTNING_INDEXER: + ggml_cuda_lightning_indexer(ctx, dst); + break; ++ case GGML_OP_SINKHORN: ++ ggml_cuda_op_sinkhorn(ctx, dst); ++ break; ++ case GGML_OP_HC_MIX: ++ ggml_cuda_op_hc_mix(ctx, dst); ++ break; ++ case GGML_OP_DSA_SCORE: ++ ggml_cuda_op_dsa_score(ctx, dst); ++ break; + default: + return false; + } +@@ -5874,6 +5885,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g + case GGML_OP_OPT_STEP_ADAMW: + case GGML_OP_OPT_STEP_SGD: + case GGML_OP_FILL: ++ case GGML_OP_SINKHORN: ++ case GGML_OP_HC_MIX: ++ case GGML_OP_DSA_SCORE: + case GGML_OP_CUMSUM: + case GGML_OP_TRI: + case GGML_OP_DIAG: +diff --git a/ggml/src/ggml-cuda/sinkhorn.cu b/ggml/src/ggml-cuda/sinkhorn.cu +new file mode 100644 +index 000000000..66add8b91 +--- /dev/null ++++ b/ggml/src/ggml-cuda/sinkhorn.cu +@@ -0,0 +1,160 @@ ++#include "sinkhorn.cuh" ++ ++// Shared by GGML_OP_SINKHORN and GGML_OP_HC_MIX: softmax along dim0, add eps, ++// then alternating dim1/dim0 normalization of the [ne0, ne1] matrix m. ++// The matrices are tiny (hyper-connection stream mixing, ne0*ne1 <= 64), so one ++// thread handles one matrix entirely in local memory. This replaces the ++// ~8 ops x 2 x n_iter graph nodes the unfused form needs per matrix. ++static __device__ void sinkhorn_matrix(float * m, const int ne0, const int ne1, const int n_iter, const float eps) { ++ const int nm = ne0*ne1; ++ ++ // softmax along dim0, per i1 ++ for (int i1 = 0; i1 < ne1; ++i1) { ++ float * c = m + i1*ne0; ++ float vmax = c[0]; ++ for (int i0 = 1; i0 < ne0; ++i0) { ++ vmax = fmaxf(vmax, c[i0]); ++ } ++ float sum = 0.0f; ++ for (int i0 = 0; i0 < ne0; ++i0) { ++ c[i0] = expf(c[i0] - vmax); ++ sum += c[i0]; ++ } ++ // reciprocal multiply to match ggml_soft_max rounding ++ const float inv = 1.0f/sum; ++ for (int i0 = 0; i0 < ne0; ++i0) { ++ c[i0] *= inv; ++ } ++ } ++ ++ // add eps to every element ++ for (int i = 0; i < nm; ++i) { ++ m[i] += eps; ++ } ++ ++ for (int it = 0; it < n_iter; ++it) { ++ // normalize along dim0, per i1 (skipped on the first pass) ++ if (it > 0) { ++ for (int i1 = 0; i1 < ne1; ++i1) { ++ float sum = eps; ++ for (int i0 = 0; i0 < ne0; ++i0) { ++ sum += m[i0 + i1*ne0]; ++ } ++ for (int i0 = 0; i0 < ne0; ++i0) { ++ m[i0 + i1*ne0] /= sum; ++ } ++ } ++ } ++ // normalize along dim1, per i0 ++ for (int i0 = 0; i0 < ne0; ++i0) { ++ float sum = eps; ++ for (int i1 = 0; i1 < ne1; ++i1) { ++ sum += m[i0 + i1*ne0]; ++ } ++ for (int i1 = 0; i1 < ne1; ++i1) { ++ m[i0 + i1*ne0] /= sum; ++ } ++ } ++ } ++} ++ ++static __global__ void sinkhorn_f32(const float * x, float * dst, ++ const int ne0, const int ne1, const int64_t nb, const int n_iter, const float eps) { ++ const int64_t b = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; ++ if (b >= nb) { ++ return; ++ } ++ ++ const int nm = ne0*ne1; ++ ++ float m[64]; ++ ++ const float * sp = x + b*nm; ++ float * dp = dst + b*nm; ++ ++ for (int i = 0; i < nm; ++i) { ++ m[i] = sp[i]; ++ } ++ ++ sinkhorn_matrix(m, ne0, ne1, n_iter, eps); ++ ++ for (int i = 0; i < nm; ++i) { ++ dp[i] = m[i]; ++ } ++} ++ ++static __global__ void hc_mix_f32(const float * x, const float * scale, const float * base, float * dst, ++ const int hc, const int64_t nt, const int n_iter, const float eps) { ++ const int64_t t = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; ++ if (t >= nt) { ++ return; ++ } ++ ++ const int n0 = 2*hc + hc*hc; ++ ++ float m[64]; ++ ++ const float * sp = x + t*n0; ++ float * dp = dst + t*n0; ++ ++ for (int i = 0; i < hc; ++i) { ++ dp[i] = 1.0f/(1.0f + expf(-(sp[i]*scale[0] + base[i]))) + eps; // pre ++ } ++ for (int i = hc; i < 2*hc; ++i) { ++ dp[i] = 2.0f/(1.0f + expf(-(sp[i]*scale[1] + base[i]))); // post ++ } ++ for (int i = 2*hc; i < n0; ++i) { ++ m[i - 2*hc] = sp[i]*scale[2] + base[i]; // comb ++ } ++ ++ sinkhorn_matrix(m, hc, hc, n_iter, eps); ++ ++ for (int i = 0; i < hc*hc; ++i) { ++ dp[2*hc + i] = m[i]; ++ } ++} ++ ++void ggml_cuda_op_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; ++ ++ GGML_ASSERT(src0->type == GGML_TYPE_F32); ++ GGML_ASSERT( dst->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_is_contiguous(src0)); ++ GGML_ASSERT(src0->ne[0]*src0->ne[1] <= 64); ++ ++ const int n_iter = ggml_get_op_params_i32(dst, 0); ++ const float eps = ggml_get_op_params_f32(dst, 1); ++ ++ const int64_t nb = src0->ne[2]*src0->ne[3]; ++ ++ const int64_t num_blocks = (nb + CUDA_SINKHORN_BLOCK_SIZE - 1)/CUDA_SINKHORN_BLOCK_SIZE; ++ ++ sinkhorn_f32<<>>( ++ (const float *) src0->data, (float *) dst->data, ++ (int) src0->ne[0], (int) src0->ne[1], nb, n_iter, eps); ++} ++ ++void ggml_cuda_op_hc_mix(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; ++ const ggml_tensor * scale = dst->src[1]; ++ const ggml_tensor * base = dst->src[2]; ++ ++ GGML_ASSERT(src0->type == GGML_TYPE_F32); ++ GGML_ASSERT( dst->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_is_contiguous(src0)); ++ ++ const int hc = ggml_get_op_params_i32(dst, 0); ++ const int n_iter = ggml_get_op_params_i32(dst, 1); ++ const float eps = ggml_get_op_params_f32(dst, 2); ++ ++ GGML_ASSERT(src0->ne[0] == 2*hc + hc*hc); ++ GGML_ASSERT(hc*hc <= 64); ++ ++ const int64_t nt = src0->ne[1]*src0->ne[2]*src0->ne[3]; ++ ++ const int64_t num_blocks = (nt + CUDA_SINKHORN_BLOCK_SIZE - 1)/CUDA_SINKHORN_BLOCK_SIZE; ++ ++ hc_mix_f32<<>>( ++ (const float *) src0->data, (const float *) scale->data, (const float *) base->data, ++ (float *) dst->data, hc, nt, n_iter, eps); ++} +diff --git a/ggml/src/ggml-cuda/sinkhorn.cuh b/ggml/src/ggml-cuda/sinkhorn.cuh +new file mode 100644 +index 000000000..8c1bcc33d +--- /dev/null ++++ b/ggml/src/ggml-cuda/sinkhorn.cuh +@@ -0,0 +1,6 @@ ++#include "common.cuh" ++ ++#define CUDA_SINKHORN_BLOCK_SIZE 128 ++ ++void ggml_cuda_op_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++void ggml_cuda_op_hc_mix(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu +index b0f609990..ea30a0204 100644 +--- a/ggml/src/ggml-cuda/top-k.cu ++++ b/ggml/src/ggml-cuda/top-k.cu +@@ -46,6 +46,94 @@ using namespace cub; + # endif + #endif // GGML_CUDA_TOP_K_NROWS_THRESHOLD + ++// map float bits to a monotonically ordered uint (larger float <=> larger uint) ++static __device__ __forceinline__ uint32_t top_k_float_to_ordered(const float f) { ++ const uint32_t u = __float_as_uint(f); ++ return (u & 0x80000000u) ? ~u : (u | 0x80000000u); ++} ++ ++// radix-select top-k, one block per row: 4 rounds of 256-bin byte histograms ++// (MSB first) narrow down the exact bit pattern of the k-th largest value, then ++// a single compaction pass collects the winning indices. O(n) passes instead of ++// a full O(n log n) sort -- at [102400, 512] this is ~5 reads of the row data ++// vs a full segmented radix sort. Output order is unspecified, which matches ++// the GGML_OP_TOP_K contract. ++#define CUDA_TOP_K_RADIX_BLOCK_SIZE 512 ++ ++static __global__ void top_k_radix_select_f32(const float * src, int * dst, const int ncols, const int k) { ++ const float * x = src + (size_t) blockIdx.x*ncols; ++ int * out = dst + (size_t) blockIdx.x*k; ++ ++ __shared__ int hist[256]; ++ __shared__ uint32_t s_prefix; ++ __shared__ int s_remaining; // elements of the current prefix group still needed ++ __shared__ int s_pos_gt, s_pos_eq; ++ ++ const int tid = threadIdx.x; ++ if (tid == 0) { ++ s_prefix = 0; ++ s_remaining = k; ++ } ++ ++ for (int round = 0; round < 4; ++round) { ++ const int shift = (3 - round)*8; ++ ++ for (int i = tid; i < 256; i += blockDim.x) { ++ hist[i] = 0; ++ } ++ __syncthreads(); ++ ++ const uint32_t prefix = s_prefix; ++ const uint32_t prefix_mask = round == 0 ? 0 : 0xFFFFFFFFu << (shift + 8); ++ for (int i = tid; i < ncols; i += blockDim.x) { ++ const uint32_t u = top_k_float_to_ordered(x[i]); ++ if ((u & prefix_mask) == prefix) { ++ const int bin = (u >> shift) & 255; ++ // warp-aggregate equal bins before hitting shared memory: masked-out ++ // rows produce millions of identical -inf entries otherwise ++ const unsigned int peers = __match_any_sync(__activemask(), bin); ++ if ((__ffs(peers) - 1) == (int) (threadIdx.x % WARP_SIZE)) { ++ atomicAdd(&hist[bin], __popc(peers)); ++ } ++ } ++ } ++ __syncthreads(); ++ ++ if (tid == 0) { ++ int rem = s_remaining; ++ int b = 255; ++ while (b > 0 && hist[b] < rem) { ++ rem -= hist[b]; ++ --b; ++ } ++ s_prefix = prefix | ((uint32_t) b << shift); ++ s_remaining = rem; ++ } ++ __syncthreads(); ++ } ++ ++ const uint32_t T = s_prefix; // ordered bits of the k-th largest value ++ const int n_gt = k - s_remaining; // elements strictly greater than T ++ ++ if (tid == 0) { ++ s_pos_gt = 0; ++ s_pos_eq = 0; ++ } ++ __syncthreads(); ++ ++ for (int i = tid; i < ncols; i += blockDim.x) { ++ const uint32_t u = top_k_float_to_ordered(x[i]); ++ if (u > T) { ++ out[atomicAdd(&s_pos_gt, 1)] = i; ++ } else if (u == T) { ++ const int slot = atomicAdd(&s_pos_eq, 1); ++ if (slot < k - n_gt) { ++ out[n_gt + slot] = i; ++ } ++ } ++ } ++} ++ + #ifdef CUB_TOP_K_AVAILABLE + + static void top_k_cub(ggml_cuda_pool & pool, +@@ -296,6 +384,15 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + + const int device = ggml_cuda_get_device(); + ++ // large batched rows: radix select, O(n) instead of a full sort (order is ++ // unspecified). single rows stay on the sort paths -- one block per row ++ // leaves the GPU idle at nrows == 1 ++ if (ncols >= 4096 && nrows > 1) { ++ top_k_radix_select_f32<<>>( ++ src0_d, dst_d, ncols, k); ++ return; ++ } ++ + #ifdef CUB_TOP_K_AVAILABLE + // a single row always uses DeviceTopK if available + const bool bitonic_short = nrows > 1 && ncols <= GGML_CUDA_TOP_K_NCOLS_THRESHOLD_BITONIC; +diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c +index 01dc8d05c..408d39f99 100644 +--- a/ggml/src/ggml.c ++++ b/ggml/src/ggml.c +@@ -1077,6 +1077,9 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + "LEAKY_RELU", + "TRI", + "FILL", ++ "SINKHORN", ++ "HC_MIX", ++ "DSA_SCORE", + + "FLASH_ATTN_EXT", + "FLASH_ATTN_BACK", +@@ -1112,7 +1115,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + "GLU", + }; + +-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); ++static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); + + static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "none", +@@ -1192,6 +1195,9 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "leaky_relu(x)", + "tri(x)", + "fill(x, c)", ++ "sinkhorn(x)", ++ "hc_mix(x)", ++ "dsa_score(x)", + + "flash_attn_ext(x)", + "flash_attn_back(x)", +@@ -1227,7 +1233,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "glu(x)", + }; + +-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); ++static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); + + static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); + +@@ -3355,8 +3361,8 @@ static inline bool ggml_can_mul_mat(const struct ggml_tensor * t0, const struct + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return (t0->ne[0] == t1->ne[0]) && +- (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable +- (t1->ne[3]%t0->ne[3] == 0); ++ (t0->ne[2] == 0 || t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable ++ (t0->ne[3] == 0 || t1->ne[3]%t0->ne[3] == 0); + } + + struct ggml_tensor * ggml_mul_mat( +@@ -3442,8 +3448,8 @@ static inline bool ggml_can_out_prod(const struct ggml_tensor * t0, const struct + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return (t0->ne[1] == t1->ne[1]) && +- (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable +- (t1->ne[3]%t0->ne[3] == 0); ++ (t0->ne[2] == 0 || t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable ++ (t0->ne[3] == 0 || t1->ne[3]%t0->ne[3] == 0); + } + + struct ggml_tensor * ggml_out_prod( +@@ -5440,6 +5446,93 @@ struct ggml_tensor * ggml_fill_inplace( + return ggml_fill_impl(ctx, a, c, true); + } + ++// ggml_sinkhorn ++ ++struct ggml_tensor * ggml_sinkhorn( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ int n_iter, ++ float eps) { ++ GGML_ASSERT(a->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_is_contiguous(a)); ++ GGML_ASSERT(n_iter >= 1); ++ // the per-batch matrix is kept in registers/stack by the backends ++ GGML_ASSERT(a->ne[0]*a->ne[1] <= 64); ++ ++ struct ggml_tensor * result = ggml_dup_tensor(ctx, a); ++ ++ ggml_set_op_params_i32(result, 0, n_iter); ++ ggml_set_op_params_f32(result, 1, eps); ++ ++ result->op = GGML_OP_SINKHORN; ++ result->src[0] = a; ++ ++ return result; ++} ++ ++// ggml_hc_mix ++ ++struct ggml_tensor * ggml_hc_mix( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * scale, ++ struct ggml_tensor * base, ++ int hc, ++ int n_iter, ++ float eps) { ++ GGML_ASSERT(a->type == GGML_TYPE_F32 && scale->type == GGML_TYPE_F32 && base->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_is_contiguous(a)); ++ GGML_ASSERT(a->ne[0] == 2*hc + hc*hc); ++ GGML_ASSERT(ggml_nelements(scale) == 3); ++ GGML_ASSERT(ggml_nelements(base) == a->ne[0]); ++ GGML_ASSERT(n_iter >= 1); ++ GGML_ASSERT(hc*hc <= 64); ++ ++ struct ggml_tensor * result = ggml_dup_tensor(ctx, a); ++ ++ ggml_set_op_params_i32(result, 0, hc); ++ ggml_set_op_params_i32(result, 1, n_iter); ++ ggml_set_op_params_f32(result, 2, eps); ++ ++ result->op = GGML_OP_HC_MIX; ++ result->src[0] = a; ++ result->src[1] = scale; ++ result->src[2] = base; ++ ++ return result; ++} ++ ++// ggml_dsa_score ++ ++struct ggml_tensor * ggml_dsa_score( ++ struct ggml_context * ctx, ++ struct ggml_tensor * ik, ++ struct ggml_tensor * q, ++ struct ggml_tensor * w, ++ struct ggml_tensor * mask) { ++ GGML_ASSERT(ik->type == GGML_TYPE_F16 || ik->type == GGML_TYPE_F32); ++ GGML_ASSERT(q->type == GGML_TYPE_F32 && ggml_is_contiguous(q)); ++ GGML_ASSERT(w->type == GGML_TYPE_F32 && ggml_is_contiguous(w)); ++ GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_is_contiguous(mask)); ++ ++ GGML_ASSERT(ik->ne[0] == q->ne[0]); // d ++ GGML_ASSERT(w->ne[0] == q->ne[1]); // h ++ GGML_ASSERT(w->ne[1] == q->ne[2]); // nt ++ GGML_ASSERT(mask->ne[0] == ik->ne[1]); // n_kv ++ GGML_ASSERT(mask->ne[1] >= q->ne[2]); ++ ++ struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ik->ne[1], q->ne[2]); ++ ++ result->op = GGML_OP_DSA_SCORE; ++ result->src[0] = ik; ++ result->src[1] = q; ++ result->src[2] = w; ++ result->src[3] = mask; ++ ++ return result; ++} ++ + // ggml_argsort + + struct ggml_tensor * ggml_argsort( +diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py +index a39d046fe..2dd64902f 100644 +--- a/gguf-py/gguf/constants.py ++++ b/gguf-py/gguf/constants.py +@@ -224,6 +224,9 @@ class Keys: + ROPE_PATTERN = "{arch}.attention.rope_pattern" + VALUE_EXPERT_COUNT = "{arch}.attention.value_expert_count" + VALUE_EXPERT_USED_COUNT = "{arch}.attention.value_expert_used_count" ++ SINK_COUNT = "{arch}.attention.sink_count" ++ CONV_KERNEL_SIZE = "{arch}.attention.conv_kernel_size" ++ LAYER_SWA_WINDOW = "{arch}.attention.layer_swa_window" + + class Indexer: + HEAD_COUNT = "{arch}.attention.indexer.head_count" +@@ -674,6 +677,7 @@ class MODEL_ARCH(IntEnum): + MINIMAXM3 = auto() + RND1 = auto() + PANGU_EMBED = auto() ++ OPENPANGU_V2 = auto() + MISTRAL3 = auto() + EAGLE3 = auto() + DFLASH = auto() +@@ -887,6 +891,12 @@ class MODEL_TENSOR(IntEnum): + PLE_NORM_QUERY = auto() # qwen4exp + PLE_NORM_CONV = auto() # qwen4exp + PLE_CONV1D = auto() # qwen4exp ++ ATTN_Q_A_CONV = auto() # openpangu-v2 MoME short conv ++ ATTN_KV_A_CONV = auto() ++ ATTN_O_CONV = auto() ++ ATTN_SINK_KV = auto() # openpangu-v2 learned param-sink compressed kv ++ ATTN_SINK_K_PE = auto() ++ BLOCK_POST_NORM = auto() # openpangu-v2 mHC stream post-norm + ATTN_COMPRESSOR_WKV = auto() + ATTN_COMPRESSOR_WGATE = auto() + ATTN_COMPRESSOR_APE = auto() +@@ -1479,6 +1489,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { + MODEL_ARCH.COGVLM: "cogvlm", + MODEL_ARCH.RND1: "rnd1", + MODEL_ARCH.PANGU_EMBED: "pangu-embedded", ++ MODEL_ARCH.OPENPANGU_V2: "openpangu-v2", + MODEL_ARCH.MISTRAL3: "mistral3", + MODEL_ARCH.EAGLE3: "eagle3", + MODEL_ARCH.DFLASH: "dflash", +@@ -1690,6 +1701,12 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { + MODEL_TENSOR.PLE_NORM_QUERY: "blk.{bid}.ple_norm_query", # qwen4exp + MODEL_TENSOR.PLE_NORM_CONV: "blk.{bid}.ple_norm_conv", # qwen4exp + MODEL_TENSOR.PLE_CONV1D: "blk.{bid}.ple_conv1d", # qwen4exp ++ MODEL_TENSOR.ATTN_Q_A_CONV: "blk.{bid}.attn_q_a_conv", ++ MODEL_TENSOR.ATTN_KV_A_CONV: "blk.{bid}.attn_kv_a_conv", ++ MODEL_TENSOR.ATTN_O_CONV: "blk.{bid}.attn_o_conv", ++ MODEL_TENSOR.ATTN_SINK_KV: "blk.{bid}.attn_sink_kv", ++ MODEL_TENSOR.ATTN_SINK_K_PE: "blk.{bid}.attn_sink_k_pe", ++ MODEL_TENSOR.BLOCK_POST_NORM: "blk.{bid}.block_post_norm", + MODEL_TENSOR.ATTN_COMPRESSOR_WKV: "blk.{bid}.attn_compressor_kv", + MODEL_TENSOR.ATTN_COMPRESSOR_WGATE: "blk.{bid}.attn_compressor_gate", + MODEL_TENSOR.ATTN_COMPRESSOR_APE: "blk.{bid}.attn_compressor_ape", +@@ -4201,6 +4218,59 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], ++ MODEL_ARCH.OPENPANGU_V2: [ ++ MODEL_TENSOR.TOKEN_EMBD, ++ MODEL_TENSOR.OUTPUT_NORM, ++ MODEL_TENSOR.OUTPUT, ++ MODEL_TENSOR.HC_HEAD_FN, ++ MODEL_TENSOR.HC_HEAD_BASE, ++ MODEL_TENSOR.HC_HEAD_SCALE, ++ MODEL_TENSOR.ATTN_NORM, ++ MODEL_TENSOR.ATTN_POST_NORM, ++ MODEL_TENSOR.BLOCK_POST_NORM, ++ MODEL_TENSOR.HC_ATTN_FN, ++ MODEL_TENSOR.HC_ATTN_BASE, ++ MODEL_TENSOR.HC_ATTN_SCALE, ++ MODEL_TENSOR.HC_FFN_FN, ++ MODEL_TENSOR.HC_FFN_BASE, ++ MODEL_TENSOR.HC_FFN_SCALE, ++ MODEL_TENSOR.ATTN_Q_A, ++ MODEL_TENSOR.ATTN_Q_A_NORM, ++ MODEL_TENSOR.ATTN_Q_B, ++ MODEL_TENSOR.ATTN_Q_A_CONV, ++ MODEL_TENSOR.ATTN_KV_A_MQA, ++ MODEL_TENSOR.ATTN_KV_A_NORM, ++ MODEL_TENSOR.ATTN_KV_A_CONV, ++ MODEL_TENSOR.ATTN_K_B, ++ MODEL_TENSOR.ATTN_V_B, ++ MODEL_TENSOR.ATTN_O_CONV, ++ MODEL_TENSOR.ATTN_OUT, ++ MODEL_TENSOR.ATTN_SINK_KV, ++ MODEL_TENSOR.ATTN_SINK_K_PE, ++ MODEL_TENSOR.INDEXER_K_NORM, ++ MODEL_TENSOR.INDEXER_PROJ, ++ MODEL_TENSOR.INDEXER_ATTN_K, ++ MODEL_TENSOR.INDEXER_ATTN_Q_B, ++ MODEL_TENSOR.FFN_GATE_INP, ++ MODEL_TENSOR.FFN_EXP_PROBS_B, ++ MODEL_TENSOR.FFN_NORM, ++ MODEL_TENSOR.FFN_POST_NORM, ++ MODEL_TENSOR.FFN_GATE, ++ MODEL_TENSOR.FFN_DOWN, ++ MODEL_TENSOR.FFN_UP, ++ 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.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, + MODEL_TENSOR.OUTPUT_NORM, +diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py +index cc2bd8e7b..d1bc15bba 100644 +--- a/gguf-py/gguf/gguf_writer.py ++++ b/gguf-py/gguf/gguf_writer.py +@@ -837,6 +837,15 @@ class GGUFWriter: + def add_indexer_kpool_select_tail(self, value: bool) -> None: + self.add_bool(Keys.Attention.Indexer.KPOOL_SELECT_TAIL.format(arch=self.arch), value) + ++ def add_sink_count(self, count: int) -> None: ++ self.add_uint32(Keys.Attention.SINK_COUNT.format(arch=self.arch), count) ++ ++ def add_conv_kernel_size(self, size: int) -> None: ++ self.add_uint32(Keys.Attention.CONV_KERNEL_SIZE.format(arch=self.arch), size) ++ ++ def add_layer_swa_window(self, values: Sequence[int]) -> None: ++ self.add_array(Keys.Attention.LAYER_SWA_WINDOW.format(arch=self.arch), values) ++ + def add_max_alibi_bias(self, bias: float) -> None: + self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias) + +diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp +index d1dd4cbae..e57d7c348 100644 +--- a/src/llama-arch.cpp ++++ b/src/llama-arch.cpp +@@ -84,6 +84,7 @@ static const std::map LLM_ARCH_NAMES = { + { LLM_ARCH_DEEPSEEK2OCR, "deepseek2-ocr" }, + { LLM_ARCH_DEEPSEEK32, "deepseek32" }, + { LLM_ARCH_DEEPSEEK4, "deepseek4" }, ++ { LLM_ARCH_OPENPANGU_V2, "openpangu-v2" }, + { LLM_ARCH_CHATGLM, "chatglm" }, + { LLM_ARCH_GLM4, "glm4" }, + { LLM_ARCH_GLM4_MOE, "glm4moe" }, +@@ -304,6 +305,9 @@ static const std::map LLM_KV_NAMES = { + { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL, "%s.attention.indexer.kpool" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, "%s.attention.indexer.kpool_select_tail" }, ++ { LLM_KV_ATTENTION_SINK_COUNT, "%s.attention.sink_count" }, ++ { LLM_KV_ATTENTION_CONV_KERNEL_SIZE, "%s.attention.conv_kernel_size" }, ++ { LLM_KV_ATTENTION_LAYER_SWA_WINDOW, "%s.attention.layer_swa_window" }, + { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" }, + { LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" }, + { LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" }, +@@ -611,6 +615,12 @@ static const std::map LLM_TENSOR_NAMES = { + { LLM_TENSOR_PLE_NORM_QUERY, "blk.%d.ple_norm_query" }, + { LLM_TENSOR_PLE_NORM_CONV, "blk.%d.ple_norm_conv" }, + { LLM_TENSOR_PLE_CONV1D, "blk.%d.ple_conv1d" }, ++ { LLM_TENSOR_ATTN_Q_A_CONV, "blk.%d.attn_q_a_conv" }, ++ { LLM_TENSOR_ATTN_KV_A_CONV, "blk.%d.attn_kv_a_conv" }, ++ { LLM_TENSOR_ATTN_O_CONV, "blk.%d.attn_o_conv" }, ++ { LLM_TENSOR_ATTN_SINK_KV, "blk.%d.attn_sink_kv" }, ++ { LLM_TENSOR_ATTN_SINK_K_PE, "blk.%d.attn_sink_k_pe" }, ++ { LLM_TENSOR_BLOCK_POST_NORM, "blk.%d.block_post_norm" }, + { LLM_TENSOR_ATTN_COMPRESSOR_WKV, "blk.%d.attn_compressor_kv" }, + { LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "blk.%d.attn_compressor_gate" }, + { LLM_TENSOR_ATTN_COMPRESSOR_APE, "blk.%d.attn_compressor_ape" }, +@@ -879,6 +889,12 @@ static const std::map LLM_TENSOR_INFOS = { + {LLM_TENSOR_PLE_NORM_QUERY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_NORM_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_CONV1D, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, ++ {LLM_TENSOR_ATTN_Q_A_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, ++ {LLM_TENSOR_ATTN_KV_A_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, ++ {LLM_TENSOR_ATTN_O_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, ++ {LLM_TENSOR_ATTN_SINK_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, ++ {LLM_TENSOR_ATTN_SINK_K_PE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, ++ {LLM_TENSOR_BLOCK_POST_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ATTN_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, +@@ -1212,6 +1228,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { + case LLM_ARCH_QWEN4EXP: + case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_MINIMAX_01: ++ case LLM_ARCH_OPENPANGU_V2: + return true; + default: + return false; +diff --git a/src/llama-arch.h b/src/llama-arch.h +index 4e8e53421..9306fe562 100644 +--- a/src/llama-arch.h ++++ b/src/llama-arch.h +@@ -89,6 +89,7 @@ enum llm_arch { + LLM_ARCH_DEEPSEEK2OCR, + LLM_ARCH_DEEPSEEK32, + LLM_ARCH_DEEPSEEK4, ++ LLM_ARCH_OPENPANGU_V2, + LLM_ARCH_CHATGLM, + LLM_ARCH_GLM4, + LLM_ARCH_GLM4_MOE, +@@ -309,6 +310,9 @@ enum llm_kv { + LLM_KV_ATTENTION_INDEXER_TYPES, + LLM_KV_ATTENTION_INDEXER_KPOOL, + LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, ++ LLM_KV_ATTENTION_SINK_COUNT, ++ LLM_KV_ATTENTION_CONV_KERNEL_SIZE, ++ LLM_KV_ATTENTION_LAYER_SWA_WINDOW, + LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, + LLM_KV_ATTENTION_OUTPUT_LORA_RANK, + LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, +@@ -655,6 +659,12 @@ enum llm_tensor { + LLM_TENSOR_PLE_NORM_QUERY, // qwen4exp + LLM_TENSOR_PLE_NORM_CONV, // qwen4exp + LLM_TENSOR_PLE_CONV1D, // qwen4exp ++ LLM_TENSOR_ATTN_Q_A_CONV, // openpangu-v2 MoME short conv ++ LLM_TENSOR_ATTN_KV_A_CONV, // openpangu-v2 ++ LLM_TENSOR_ATTN_O_CONV, // openpangu-v2 ++ LLM_TENSOR_ATTN_SINK_KV, // openpangu-v2 learned param-sink compressed kv ++ LLM_TENSOR_ATTN_SINK_K_PE, // openpangu-v2 ++ LLM_TENSOR_BLOCK_POST_NORM,// openpangu-v2 mHC stream post-norm + LLM_TENSOR_ATTN_COMPRESSOR_WKV, + LLM_TENSOR_ATTN_COMPRESSOR_WGATE, + LLM_TENSOR_ATTN_COMPRESSOR_APE, +diff --git a/src/llama-context.cpp b/src/llama-context.cpp +index 9ee64b566..5177c40c5 100644 +--- a/src/llama-context.cpp ++++ b/src/llama-context.cpp +@@ -2506,7 +2506,10 @@ void llama_context::output_reorder() { + + uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { + uint32_t res; +- if (model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_GLM5_NEXT) { ++ if (model.arch == LLM_ARCH_OPENPANGU_V2) { ++ // mHC (per-layer sinkhorn over many iterations) creates a very large graph ++ res = std::max(n_tokens * 64, 160u * model.n_tensors()); ++ } else if (model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_GLM5_NEXT) { + // the n_tokens*40 budget below is exhausted at ubatch 3840 + res = std::max(n_tokens * 160, 64u * model.n_tensors()); + } else if (model.arch == LLM_ARCH_HRM_TEXT) { +diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp +index 6cadda971..913b6a68a 100644 +--- a/src/llama-graph.cpp ++++ b/src/llama-graph.cpp +@@ -527,7 +527,7 @@ void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { + if (self_k_idxs && self_k_idxs->buffer) { + mctx->set_input_k_idxs(self_k_idxs, ubatch); + } +- if (self_v_idxs && self_v_idxs->buffer) { ++ if (self_v_idxs && self_v_idxs->buffer) { // not built for K-only MLA caches + mctx->set_input_v_idxs(self_v_idxs, ubatch); + } + +@@ -1148,7 +1148,9 @@ void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) { + + void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { + mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); +- mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); ++ if (inp_attn->self_v_idxs) { // not built for K-only MLA caches ++ mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); ++ } + + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + +@@ -1243,7 +1245,9 @@ void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) { + // base tensors may not be allocated if there are no non-SWA attention layers + if (inp_attn->self_k_idxs && inp_attn->self_k_idxs->buffer) { + attn_ctx->get_base()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); +- attn_ctx->get_base()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); ++ if (inp_attn->self_v_idxs) { // not built for K-only MLA caches ++ attn_ctx->get_base()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); ++ } + } + + if (inp_attn->self_kq_mask && inp_attn->self_kq_mask->buffer) { +@@ -1253,7 +1257,9 @@ void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) { + // swa tensors may not be allocated if there are no SWA attention layers + if (inp_attn->self_k_idxs_swa && inp_attn->self_k_idxs_swa->buffer) { + attn_ctx->get_swa()->set_input_k_idxs(inp_attn->self_k_idxs_swa, ubatch); +- attn_ctx->get_swa()->set_input_v_idxs(inp_attn->self_v_idxs_swa, ubatch); ++ if (inp_attn->self_v_idxs_swa) { // not built for K-only MLA caches ++ attn_ctx->get_swa()->set_input_v_idxs(inp_attn->self_v_idxs_swa, ubatch); ++ } + } + + if (inp_attn->self_kq_mask_swa && inp_attn->self_kq_mask_swa->buffer) { +diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp +index b83f45ba9..f025555ee 100644 +--- a/src/llama-hparams.cpp ++++ b/src/llama-hparams.cpp +@@ -216,6 +216,12 @@ uint32_t llama_hparams::n_embd_r() const { + return n_embd * (n_shortconv_l_cache - 1); + } + ++ if (openpangu_conv_k != 0) { ++ // openPangu MoME: three causal convs (q_a, kv-compressed, attn-out) per layer ++ const uint32_t o_dim = n_head() * n_embd_head_v_mla(); ++ return (openpangu_conv_k - 1) * (n_lora_q + n_lora_kv + o_dim); ++ } ++ + if (n_embd_head_kda != 0) { + // for Kimi KDA layers + // Conv state for Q, K, V: 3 * (d_conv - 1) * n_head * head_dim +diff --git a/src/llama-hparams.h b/src/llama-hparams.h +index 922164a54..a316d2521 100644 +--- a/src/llama-hparams.h ++++ b/src/llama-hparams.h +@@ -351,6 +351,14 @@ struct llama_hparams { + // PLE conv history rows: (kernel - 1) * ngram_size; 0 without a PLE module + uint32_t ple_conv_state() const; + ++ // openPangu-2.0 (openpangu-v2) ++ uint32_t openpangu_n_sink = 0; // learned param-sink KV entries per layer ++ uint32_t openpangu_conv_k = 0; // MoME causal conv kernel width ++ ++ // reserved always-visible rows at the head of each K cache stream slice ++ // (learned attention sinks live there; K-only/MLA caches only) ++ uint32_t n_k_sink_prefix = 0; ++ + // qwen3vl deepstack + // When parsed from GGUF, this implies the first N layers consume the first + // N deepstack embeddings. Use deepstack_mapping_arr if you need a more +diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp +index b7dfadef4..baf89e59a 100644 +--- a/src/llama-kv-cache.cpp ++++ b/src/llama-kv-cache.cpp +@@ -230,7 +230,14 @@ llama_kv_cache::llama_kv_cache( + const bool has_k = true; + const bool has_v = !is_mla; + +- ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr; ++ // reserved always-visible rows at the head of each K stream slice (attention sinks) ++ const uint32_t n_k_prefix = hparams.n_k_sink_prefix; ++ if (n_k_prefix > 0) { ++ GGML_ASSERT(n_stream == 1 && "K sink prefix requires a unified KV cache"); ++ GGML_ASSERT(!has_v && "K sink prefix is only supported for K-only (MLA) caches"); ++ } ++ ++ ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size + n_k_prefix, n_stream) : nullptr; + ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr; + + has_k && ggml_format_name(k, "cache_%sk_l%d", name_tag, il); +@@ -240,7 +247,8 @@ llama_kv_cache::llama_kv_cache( + std::vector v_stream; + + for (uint32_t s = 0; s < n_stream; ++s) { +- k_stream.push_back(has_k ? ggml_view_2d(ctx, k, n_embd_k_gqa, kv_size, k->nb[1], s*k->nb[2]) : nullptr); ++ // the per-stream cell region starts after the reserved sink prefix rows ++ k_stream.push_back(has_k ? ggml_view_2d(ctx, k, n_embd_k_gqa, kv_size, k->nb[1], s*k->nb[2] + n_k_prefix*k->nb[1]) : nullptr); + v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr); + } + +@@ -335,6 +343,12 @@ llama_kv_cache::llama_kv_cache( + attn_rot_k = true; + } + ++ // openPangu-v2's custom MLA+sink attention does not apply the rotation, ++ // so do not create rotation inputs its graph would leave unused ++ if (model.arch == LLM_ARCH_OPENPANGU_V2) { ++ attn_rot_k = false; ++ } ++ + const bool attn_rot_v = + !attn_rot_disable && + n_embd_head_v_all > 0 && +@@ -1300,19 +1314,23 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k + + auto * k = layers[ikv].k; + +- const uint64_t kv_size = get_size(); + const uint64_t n_embd_k_gqa = k->ne[0]; + ++ // per-stream slice: [n_k_sink_prefix reserved rows | kv_size cell rows] ++ const uint64_t slice_size = k->ne[1]; ++ const uint32_t n_k_prefix = hparams.n_k_sink_prefix; ++ + assert(n_embd_k_gqa == hparams.n_embd_k_gqa(il)); + + const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; + ++ // the view starts at the sink prefix so attention sees [sinks | cells] contiguously + return ggml_view_4d(ctx, k, +- hparams.n_embd_head_k(il), hparams.n_head_kv(il), n_kv, ns, ++ hparams.n_embd_head_k(il), hparams.n_head_kv(il), n_kv + n_k_prefix, ns, + ggml_row_size(k->type, hparams.n_embd_head_k(il)), + ggml_row_size(k->type, n_embd_k_gqa), +- ggml_row_size(k->type, n_embd_k_gqa*kv_size), +- ggml_row_size(k->type, n_embd_k_gqa*kv_size)*sinfo.s0); ++ ggml_row_size(k->type, n_embd_k_gqa*slice_size), ++ ggml_row_size(k->type, n_embd_k_gqa*slice_size)*sinfo.s0); + } + + ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { +@@ -1366,9 +1384,13 @@ ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggm + + k_cur = ggml_view_2d(ctx, k_cur, n_embd_gqa, n_tokens, k_cur->nb[2], 0); + +- const int64_t n_stream = k->ne[2]; ++ const int64_t n_stream = k->ne[2]; ++ const uint32_t n_k_prefix = hparams.n_k_sink_prefix; + +- if (n_stream > 1) { ++ if (n_k_prefix > 0) { ++ // cell rows start after the reserved sink prefix (n_stream == 1, asserted at init) ++ k = ggml_view_2d(ctx, k, n_embd_gqa, get_size(), k->nb[1], n_k_prefix*k->nb[1]); ++ } else if (n_stream > 1) { + const int64_t kv_size = get_size(); + + assert(n_embd_gqa == k->ne[0]); +@@ -1450,6 +1472,11 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama + } + + ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { ++ if (hparams.is_mla()) { ++ // K-only MLA cache: there are no V tensors, so no graph consumes v_idxs ++ return nullptr; ++ } ++ + const uint32_t n_tokens = ubatch.n_tokens; + + ggml_tensor * v_idxs; +@@ -2058,6 +2085,8 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co + n_rot, n_head_kv, get_size()*n_stream, + ggml_row_size(layer.k->type, n_embd_head_k), + ggml_row_size(layer.k->type, n_embd_k_gqa), ++ // cell rows start after the reserved sink prefix (n_stream == 1 when prefix > 0) ++ hparams.n_k_sink_prefix*ggml_row_size(layer.k->type, n_embd_k_gqa) + + ggml_row_size(layer.k->type, n_embd_nope)); + + ggml_tensor * cur = build_rope_shift(cparams, ctx, k, inp->k_shift, inp->k_rot, rope_factors, freq_base_l, freq_scale_l, il); +diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp +index 2893d819e..0627ae6b8 100644 +--- a/src/llama-memory-recurrent.cpp ++++ b/src/llama-memory-recurrent.cpp +@@ -26,7 +26,10 @@ llama_memory_recurrent::llama_memory_recurrent( + uint32_t n_seq_max, + uint32_t n_rs_seq, + const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max) { +- const int32_t n_layer = hparams.n_layer(); ++ // size to the full layer count (incl. MTP/NextN layers) so a filtered MTP ++ // context can allocate recurrent state for layers >= n_layer(); get_r_l/get_s_l ++ // index by the original layer id, and the filter leaves unused layers null. ++ const int32_t n_layer = hparams.n_layer_all; + + head = 0; + size = mem_size; +diff --git a/src/llama-model.cpp b/src/llama-model.cpp +index ded790b2b..8aa61719e 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -206,6 +206,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params + return new llama_model_dots3note(params); + case LLM_ARCH_DEEPSEEK4: + return new llama_model_deepseek4(params); ++ case LLM_ARCH_OPENPANGU_V2: ++ return new llama_model_openpangu_v2(params); + case LLM_ARCH_GLM_DSA: + return new llama_model_glm_dsa(params); + case LLM_ARCH_MISTRAL4: +@@ -3069,10 +3071,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + default: + { + // Dense MTP heads use a plain attention KV cache instead of the hybrid wrapper. ++ // openPangu-v2 also qualifies: its MTP layers skip the recurrent MoME conv. + const bool mtp_on_hybrid_qwen = + params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && + (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || +- arch == LLM_ARCH_BAILINGMOE3); ++ arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_OPENPANGU_V2); + + const bool mtp_on_hybrid_nemotron = + params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE; +@@ -3096,7 +3099,17 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + // a null filter_idx means the GGUF has no indexer tensors + llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; + const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP); +- if (arch == LLM_ARCH_FALCON_H1) { ++ if (arch == LLM_ARCH_OPENPANGU_V2 && hparams.n_layer_nextn > 0) { ++ // every layer has both attention (MLA) and a recurrent conv state; ++ // split the trunk (default context) from the MTP heads (MTP context) ++ const bool is_mtp = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP; ++ auto filt = [is_mtp, this](uint32_t il) { ++ return is_mtp ? (il >= hparams.n_layer()) : (il < hparams.n_layer()); ++ }; ++ filter_attn = filt; ++ filter_recr = filt; ++ } else if (arch == LLM_ARCH_FALCON_H1 || arch == LLM_ARCH_OPENPANGU_V2) { ++ // every layer has both attention (MLA) and a recurrent conv state + filter_attn = [&](uint32_t) { return true; }; + filter_recr = [&](uint32_t) { return true; }; + } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { +@@ -3492,6 +3505,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { + case LLM_ARCH_DEEPSEEK32: + case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_MUSE_GLIMMER: ++ case LLM_ARCH_OPENPANGU_V2: + case LLM_ARCH_PLM: + case LLM_ARCH_CHATGLM: + case LLM_ARCH_GRANITE: +diff --git a/src/llama-model.h b/src/llama-model.h +index bcfbf1e2c..ae137b181 100644 +--- a/src/llama-model.h ++++ b/src/llama-model.h +@@ -537,6 +537,15 @@ struct llama_layer { + struct ggml_tensor * hc_ffn_fn = nullptr; + struct ggml_tensor * hc_ffn_base = nullptr; + struct ggml_tensor * hc_ffn_scale = nullptr; ++ ++ // openPangu-2.0 (openpangu-v2) ++ struct ggml_tensor * attn_q_a_conv = nullptr; ++ struct ggml_tensor * attn_kv_a_conv = nullptr; ++ struct ggml_tensor * attn_o_conv = nullptr; ++ struct ggml_tensor * attn_sink_kv = nullptr; ++ struct ggml_tensor * attn_sink_k_pe = nullptr; ++ struct ggml_tensor * block_post_norm = nullptr; ++ + struct ggml_tensor * attn_comp_wkv = nullptr; + struct ggml_tensor * attn_comp_wgate = nullptr; + struct ggml_tensor * attn_comp_ape = nullptr; +diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp +index 9c6066796..33456ef2f 100644 +--- a/src/llama-quant.cpp ++++ b/src/llama-quant.cpp +@@ -351,6 +351,17 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param + quantize &= name.compare(0, 9, "decision.") != 0; + } + ++ // openPangu-2.0: keep MoME convs, mHC phi/scale/base, and param sinks unquantized. ++ // The hc_* names are shared with DeepSeek-V4, whose (larger) hc tensors do get ++ // quantized, so restrict these rules to this arch. ++ if (arch == LLM_ARCH_OPENPANGU_V2) { ++ quantize &= name.find("_conv.weight") == std::string::npos; ++ quantize &= name.find("attn_sink_") == std::string::npos; ++ quantize &= name.find("hc_attn") == std::string::npos; ++ quantize &= name.find("hc_ffn") == std::string::npos; ++ quantize &= name.find("output_hc") == std::string::npos; ++ } ++ + // do not quantize RWKV's small yet 2D weights + quantize &= name.find("time_mix_first.weight") == std::string::npos; + quantize &= name.find("time_mix_w0.weight") == std::string::npos; +diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp +index 9b5cc2e1c..b956e541e 100644 +--- a/src/llama-vocab.cpp ++++ b/src/llama-vocab.cpp +@@ -437,6 +437,7 @@ struct llm_tokenizer_bpe : llm_tokenizer { + break; + case LLAMA_VOCAB_PRE_TYPE_GPT4O: + case LLAMA_VOCAB_PRE_TYPE_MINIMAX_M2: ++ case LLAMA_VOCAB_PRE_TYPE_OPENPANGU_V2: + regex_exprs = { + // original regex from tokenizer.json + // "[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]*[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\\r\\n\\p{L}\\p{N}]?[\\p{Lu}\\p{Lt}\\p{Lm}\\p{Lo}\\p{M}]+[\\p{Ll}\\p{Lm}\\p{Lo}\\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", +@@ -2475,6 +2476,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { + tokenizer_pre == "talkie") { + pre_type = LLAMA_VOCAB_PRE_TYPE_GPT4O; + clean_spaces = false; ++ } else if ( ++ tokenizer_pre == "openpangu-v2") { ++ pre_type = LLAMA_VOCAB_PRE_TYPE_OPENPANGU_V2; ++ clean_spaces = false; + } else if ( + tokenizer_pre == "granite-embed-multi-97m") { + pre_type = LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI; +@@ -3051,6 +3056,8 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { + || t.first == "<｜end▁of▁sentence｜>" // deepseek-ocr + || t.first == "[e~[" // minimax-m2/m3 + || t.first == "<|ifm|im_end|>" // k2-horizon ++ || t.first == "<|message_end|>" // openpangu ++ || t.first == "<|pangu_text_end|>" // openpangu + ) { + special_eog_ids.insert(t.second); + if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) { +diff --git a/src/llama-vocab.h b/src/llama-vocab.h +index 898a23ae6..8222b0068 100644 +--- a/src/llama-vocab.h ++++ b/src/llama-vocab.h +@@ -70,6 +70,7 @@ enum llama_vocab_pre_type { + LLAMA_VOCAB_PRE_TYPE_UFAKZEKA = 59, + LLAMA_VOCAB_PRE_TYPE_MMBERT = 60, + LLAMA_VOCAB_PRE_TYPE_K2_HORIZON = 61, ++ LLAMA_VOCAB_PRE_TYPE_OPENPANGU_V2 = 62, + }; + + struct LLM_KV; +diff --git a/src/models/models.h b/src/models/models.h +index 88fd07d97..18d046347 100644 +--- a/src/models/models.h ++++ b/src/models/models.h +@@ -1407,6 +1407,81 @@ struct llama_model_deepseek4 : public llama_model_base { + }; + + ++struct llama_model_openpangu_v2 : public llama_model_base { ++ llama_model_openpangu_v2(const struct llama_model_params & params) : llama_model_base(params) {} ++ void load_arch_hparams(llama_model_loader & ml) override; ++ void load_arch_tensors(llama_model_loader & ml) override; ++ ++ // true when every model device supports the fused GGML_OP_SINKHORN; computed ++ // lazily on the first graph build (backends without the op, e.g. Metal, get ++ // the unfused op sequence instead) ++ bool sinkhorn_fused() const; ++ mutable int sinkhorn_fused_probe = -1; ++ ++ // shared helpers used by both the trunk graph and the MTP head graph ++ struct graph_base : public llm_graph_context { ++ graph_base(const llm_graph_params & params, bool sinkhorn_fused) ++ : llm_graph_context(params), use_fused_sinkhorn(sinkhorn_fused) {} ++ ++ // whether build_hc_sinkhorn emits the fused op or the unfused sequence ++ const bool use_fused_sinkhorn; ++ ++ // hoisted per-graph mask tensors (built once by the concrete graph ctor): ++ // [n_k_sink_prefix | n_kv] combined masks and the standalone sink block ++ ggml_tensor * kq_mask_sinked = nullptr; ++ ggml_tensor * kq_mask_sinked_swa = nullptr; ++ ggml_tensor * sink_mask_blk = nullptr; ++ ++ // sink-prefix mask block [n_k_sink_prefix, n_tokens]: real sinks visible, padding masked ++ ggml_tensor * build_sink_mask_blk(ggml_tensor * kq_mask) const; ++ ++ // write the (static) param-sink rows into the reserved K-cache prefix ++ void build_sink_write(const llama_layer & layer, ggml_tensor * k_row, int il) const; ++ ++ // mHC (manifold hyper-connections), gamma pre-folded into phi at conversion. ++ // xt/residual_t is the transposed stream state [hc, n_embd, nt], computed once ++ // per block and shared by the pre- and post-merge batched matmuls; when null ++ // (large batches) the per-stream loop form is used instead. ++ ggml_tensor * build_hc_pre(ggml_tensor * x, ggml_tensor * xt, ++ ggml_tensor * hc_fn, ggml_tensor * hc_scale, ++ ggml_tensor * hc_base, ggml_tensor ** post, ggml_tensor ** comb, int il) const; ++ ggml_tensor * build_hc_post(ggml_tensor * x, ggml_tensor * residual, ggml_tensor * residual_t, ++ ggml_tensor * post, ggml_tensor * comb, int il) const; ++ ggml_tensor * build_hc_head(ggml_tensor * x, ggml_tensor * hc_fn, ggml_tensor * hc_scale, ++ ggml_tensor * hc_base) const; ++ ggml_tensor * build_hc_weighted_sum(ggml_tensor * x, ggml_tensor * xt, ggml_tensor * weights) const; ++ ggml_tensor * build_hc_sinkhorn(ggml_tensor * comb, int il) const; ++ ++ // MoME width-k causal depthwise conv + identity residual, with recurrent decode state ++ ggml_tensor * build_mome_conv(ggml_tensor * x, ggml_tensor * conv_w, ++ ggml_tensor * conv_rs, ggml_tensor * conv_state, int64_t rs_head, ++ int64_t state_off, int64_t n_seqs, int64_t n_seq_tokens) const; ++ ++ ggml_tensor * build_attention(const llama_model & model, llm_graph_input_mem_hybrid_iswa * inp, ++ ggml_tensor * cur, ggml_tensor * inp_pos, float kq_scale, int il) const; ++ ++ // MTP-head attention: MLA + param sinks over a plain (non-hybrid) iSWA KV cache, ++ // without the recurrent MoME conv (drafts are verified, so the conv is skipped). ++ ggml_tensor * build_attention_mtp(const llama_model & model, llm_graph_input_attn_kv_iswa * inp_attn, ++ ggml_tensor * cur, ggml_tensor * inp_pos, float kq_scale, int il) const; ++ ++ // MoE + shared-expert block (used by both trunk and MTP layers) ++ ggml_tensor * build_moe_block(const llama_model & model, ggml_tensor * cur, int il) const; ++ }; ++ ++ struct graph : public graph_base { ++ graph(const llama_model & model, const llm_graph_params & params); ++ }; ++ ++ // MTP / NextN head graph for self-speculative decoding ++ struct graph_mtp : public graph_base { ++ graph_mtp(const llama_model & model, const llm_graph_params & params); ++ }; ++ ++ std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; ++}; ++ ++ + struct llama_model_deepseek2ocr : public llama_model_base { + llama_model_deepseek2ocr(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; +diff --git a/src/models/openpangu-v2.cpp b/src/models/openpangu-v2.cpp +new file mode 100644 +index 000000000..38f1a8c14 +--- /dev/null ++++ b/src/models/openpangu-v2.cpp +@@ -0,0 +1,1048 @@ ++#include "models.h" ++ ++#include "ggml-cpp.h" ++#include "llama-kv-cache.h" ++#include "llama-kv-cache-iswa.h" ++#include "llama-memory-hybrid.h" ++#include "llama-memory-hybrid-iswa.h" ++#include "llama-memory-recurrent.h" ++ ++// openPangu-2.0-Flash (openpangu-v2) ++// ++// Composition of existing infra: ++// - 4-stream mHC residual (manifold hyper-connections), same family as DeepSeek-V4. ++// The learned RMSNorm gamma is folded into the phi weight at conversion time. ++// - MLA attention (absorbed), same as DeepSeek-V2, scale 1/sqrt(qk_head_dim). ++// - MoME: width-k causal depthwise conv + identity residual on q_a, kv-compressed, and attn-out. ++// - 128 learned param-sink KV entries per layer (always visible, not indexer-scored). ++// - Sigmoid-gated MoE with expert bias + shared expert, leading dense layers. ++// - DSA lightning indexer on the non-SWA trunk layers: per-token top-k selection over ++// the cached tokens, applied as a mask on the full causal attention. The indexer key ++// rides in the tail of the cached MLA K row (no separate indexer cache). ++// ++// The MoME conv state is carried across ubatches via the recurrent memory; SWA layers use ++// the iSWA split cache. ++ ++static ggml_tensor * opv2_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) { ++ return ggml_view_1d(ctx, t, ne0, i0 * ggml_element_size(t)); ++} ++ ++static ggml_tensor * opv2_view_2d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t i0) { ++ return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], i0 * ggml_element_size(t)); ++} ++ ++static ggml_tensor * opv2_hc_affine(ggml_context * ctx, ggml_tensor * x, ggml_tensor * scale, ggml_tensor * base) { ++ x = ggml_mul(ctx, x, scale); ++ x = ggml_add(ctx, x, base); ++ return x; ++} ++ ++void llama_model_openpangu_v2::load_arch_hparams(llama_model_loader & ml) { ++ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); ++ ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); ++ ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); ++ ++ ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); ++ ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); ++ ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl, false); ++ ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl, false); ++ ++ // DSA lightning indexer (top-k token selection on the non-SWA trunk layers) ++ ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head, false); ++ ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size, false); ++ ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false); ++ ++ ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); ++ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); ++ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); ++ ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); ++ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); ++ ++ // mHC (reuse the DeepSeek-V4 hyper-connection fields) ++ ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); ++ ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); ++ ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); ++ ++ // required: the param-sink and MoME conv tensors are sized from these ++ ml.get_key(LLM_KV_ATTENTION_SINK_COUNT, hparams.openpangu_n_sink); ++ ml.get_key(LLM_KV_ATTENTION_CONV_KERNEL_SIZE, hparams.openpangu_conv_k); ++ ++ // the param sinks live in reserved rows at the head of the K cache; pad the ++ // prefix to the flash-attention stride so the total KV length stays aligned ++ hparams.n_k_sink_prefix = GGML_PAD(hparams.openpangu_n_sink, 256); ++ ++ // per-layer sliding window: SWA layers use a local window, DSA layers are global. ++ // Both read from the same compressed MLA KV cache; the iSWA split keeps the ++ // windowed layers cheap in memory at long context. ++ std::array swa_window; ++ swa_window.fill(0); ++ ml.get_arr(LLM_KV_ATTENTION_LAYER_SWA_WINDOW, swa_window, false); ++ uint32_t swa = 0; ++ for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { ++ hparams.is_swa_impl[il] = swa_window[il] != 0; ++ if (swa_window[il] != 0) { ++ swa = swa_window[il]; ++ } ++ } ++ if (swa == 0) { ++ // the graph and memory creation assume the hybrid-iswa layout ++ throw std::runtime_error("openpangu-v2 requires attention.layer_swa_window with at least one non-zero window"); ++ } ++ hparams.n_swa = swa; ++ hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; ++ ++ // MLA stores a single compressed KV head: size the KV cache to the ++ // compressed dims (kv_lora + rope for K, kv_lora for V). The real ++ // per-head attention dims come from n_embd_head_{k,v}_mla(). ++ // DSA layers append their indexer key to the same cache row; layers ++ // without an indexer (SWA, MTP) zero-pad the tail. ++ hparams.n_embd_head_k_full = hparams.n_lora_kv + hparams.n_rot() + hparams.indexer_head_size; ++ hparams.n_embd_head_v_full = hparams.n_lora_kv; ++ // SWA layers share the same compressed MLA KV geometry (used to size the iSWA SWA sub-cache) ++ hparams.n_embd_head_k_swa = hparams.n_embd_head_k_full; ++ hparams.n_embd_head_v_swa = hparams.n_embd_head_v_full; ++ hparams.n_rot_swa = hparams.n_rot(); ++ ++ type = LLM_TYPE_UNKNOWN; ++} ++ ++void llama_model_openpangu_v2::load_arch_tensors(llama_model_loader &) { ++ LLAMA_LOAD_LOCALS; ++ ++ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); ++ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla(); ++ const int64_t n_embd_head_qk_rope = hparams.n_rot(); ++ const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; ++ ++ const int64_t q_lora_rank = hparams.n_lora_q; ++ const int64_t kv_lora_rank = hparams.n_lora_kv; ++ const int64_t n_ff_exp = hparams.n_ff_exp(); ++ const int64_t n_expert_shared = hparams.n_expert_shared; ++ const int64_t hc = hparams.dsv4_hc_mult; ++ const int64_t conv_k = hparams.openpangu_conv_k; ++ const int64_t n_sink = hparams.openpangu_n_sink; ++ ++ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); ++ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); ++ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); ++ ++ // model-level mHC stream merge (pre_only); absent in split MTP-only GGUFs ++ const int mhc_flags = n_layer > 0 ? 0 : TENSOR_NOT_REQUIRED; ++ hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc * n_embd, hc}, mhc_flags); ++ hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc}, mhc_flags); ++ hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, mhc_flags); ++ ++ const int64_t hc_mix_dim = (2 + hc) * hc; ++ ++ for (int i = 0; i < (int) hparams.n_layer_all; ++i) { ++ auto & layer = layers[i]; ++ const bool is_mtp = i >= n_layer; // MTP / NextN heads have no mHC or indexer ++ ++ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); ++ layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0); ++ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); ++ layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); ++ ++ if (!is_mtp) { ++ layer.block_post_norm = create_tensor(tn(LLM_TENSOR_BLOCK_POST_NORM, "weight", i), {hc * n_embd}, TENSOR_NOT_REQUIRED); ++ ++ // mHC per-layer ++ layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc * n_embd, 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 * n_embd, 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); ++ } else { ++ // MTP / NextN head projections (shared_head/embed fall back to model.tok_embd/output) ++ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, 0); ++ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, 0); ++ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, 0); ++ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); ++ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); ++ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); ++ } ++ ++ // MLA ++ layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0); ++ layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0); ++ layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, 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_k_mla}, 0); ++ layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA,"weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, 0); ++ layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, 0); ++ layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, 0); ++ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, 0); ++ ++ // MoME convs (stored [k, dim]) ++ layer.attn_q_a_conv = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_CONV, "weight", i), {conv_k, q_lora_rank}, 0); ++ layer.attn_kv_a_conv = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_CONV, "weight", i), {conv_k, kv_lora_rank}, 0); ++ layer.attn_o_conv = create_tensor(tn(LLM_TENSOR_ATTN_O_CONV, "weight", i), {conv_k, n_head * n_embd_head_v_mla}, 0); ++ ++ // learned param sinks ++ layer.attn_sink_kv = create_tensor(tn(LLM_TENSOR_ATTN_SINK_KV, "weight", i), {kv_lora_rank, n_sink}, 0); ++ layer.attn_sink_k_pe = create_tensor(tn(LLM_TENSOR_ATTN_SINK_K_PE, "weight", i), {n_embd_head_qk_rope, n_sink}, 0); ++ ++ // DSA lightning indexer (present only on the non-SWA trunk layers) ++ const int64_t idx_h = hparams.indexer_n_head; ++ const int64_t idx_d = hparams.indexer_head_size; ++ layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, idx_h * idx_d}, TENSOR_NOT_REQUIRED); ++ layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, idx_d}, TENSOR_NOT_REQUIRED); ++ layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {idx_d}, TENSOR_NOT_REQUIRED); ++ layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, idx_h}, TENSOR_NOT_REQUIRED); ++ ++ if (i < (int) hparams.n_layer_dense_lead) { ++ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); ++ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); ++ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); ++ } else { ++ 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}, TENSOR_NOT_REQUIRED); ++ 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); ++ } ++ } ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_sink_mask_blk(ggml_tensor * kq_mask) const { ++ const int64_t n_sink = hparams.openpangu_n_sink; ++ const int64_t n_pfx = hparams.n_k_sink_prefix; ++ ++ ggml_tensor * m_vis = ggml_new_tensor_4d(ctx0, kq_mask->type, n_sink, ++ kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]); ++ m_vis = ggml_fill(ctx0, m_vis, 0.0f); // real sinks: always visible ++ ggml_tensor * blk = m_vis; ++ if (n_pfx > n_sink) { ++ ggml_tensor * m_pad = ggml_new_tensor_4d(ctx0, kq_mask->type, n_pfx - n_sink, ++ kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]); ++ m_pad = ggml_fill(ctx0, m_pad, -INFINITY); // padding rows: masked out ++ blk = ggml_concat(ctx0, m_vis, m_pad, 0); ++ } ++ return blk; ++} ++ ++void llama_model_openpangu_v2::graph_base::build_sink_write(const llama_layer & layer, ggml_tensor * k_row, int il) const { ++ const int64_t n_sink = hparams.openpangu_n_sink; ++ const int64_t n_pfx = hparams.n_k_sink_prefix; ++ const int64_t idx_d = hparams.indexer_head_size; ++ ++ // sink rows [row_width, n_pfx]: RMS-norm'd compressed KV + k_pe (not RoPE'd), ++ // zero indexer tail and zero (masked) padding rows. The values are static; ++ // rewriting the small prefix every graph avoids a load-time hook. ++ ggml_tensor * sink_in = ggml_cast(ctx0, layer.attn_sink_kv, GGML_TYPE_F32); ++ ggml_tensor * sink_kv = build_norm(sink_in, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); ++ ggml_tensor * sink_pe = ggml_cast(ctx0, layer.attn_sink_k_pe, sink_kv->type); ++ ggml_tensor * sink_k = ggml_concat(ctx0, sink_kv, sink_pe, 0); ++ if (idx_d > 0 || n_pfx > n_sink) { ++ sink_k = ggml_pad(ctx0, sink_k, idx_d, n_pfx - n_sink, 0, 0); ++ } ++ sink_k = ggml_cast(ctx0, sink_k, k_row->type); ++ ++ ggml_tensor * dst = ggml_view_2d(ctx0, k_row, k_row->ne[0], n_pfx, k_row->nb[2], 0); ++ ggml_build_forward_expand(gf, ggml_cpy(ctx0, sink_k, dst)); ++} ++ ++// The mHC stream reductions have two shapes: a batched-matmul form (few kernel ++// launches -- wins at decode, where launch overhead dominates) and a per-stream ++// loop form (no stream-transpose copies or tiny batched GEMMs -- wins at prefill). ++// The transposed stream state xt selects the form: non-null => batched. ++static const int64_t OPV2_HC_BATCHED_NT_MAX = 8; ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_hc_weighted_sum(ggml_tensor * x, ggml_tensor * xt, ggml_tensor * weights) const { ++ const int64_t hc = hparams.dsv4_hc_mult; ++ const int64_t nt = x->ne[2]; ++ ++ if (xt) { ++ // out[e,t] = sum_h xt[h,e,t] * w[h,t]: one batched matvec over the streams ++ if (!ggml_is_contiguous(weights)) { ++ weights = ggml_cont(ctx0, weights); // hc_mix output views are strided ++ } ++ ggml_tensor * wr = ggml_reshape_3d(ctx0, weights, hc, 1, nt); ++ ggml_tensor * out = ggml_mul_mat(ctx0, xt, wr); // [n_embd, 1, nt] ++ return ggml_reshape_2d(ctx0, out, n_embd, nt); ++ } ++ ++ ggml_tensor * acc = nullptr; ++ for (int64_t ih = 0; ih < hc; ++ih) { ++ ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih * x->nb[1]); ++ ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih * weights->nb[0]); ++ ggml_tensor * cur = ggml_mul(ctx0, xh, wh); ++ acc = acc ? ggml_add(ctx0, acc, cur) : cur; ++ } ++ return acc; ++} ++ ++bool llama_model_openpangu_v2::sinkhorn_fused() const { ++ if (sinkhorn_fused_probe < 0) { ++ sinkhorn_fused_probe = 1; ++ ++ for (const llama_device & ldev : devices) { ++ if (ldev.is_meta) { ++ continue; ++ } ++ ggml_backend_dev_t dev = ldev.dev; ++ ++ ggml_init_params params = { ++ /*.mem_size =*/ ggml_tensor_overhead()*16, ++ /*.mem_buffer =*/ NULL, ++ /*.no_alloc =*/ true, ++ }; ++ ggml_context_ptr ctx { ggml_init(params) }; ++ GGML_ASSERT(ctx); ++ ++ ggml_backend_buffer_type_t buft = ggml_backend_dev_buffer_type(dev); ++ ggml_backend_buffer_ptr buf { ggml_backend_buft_alloc_buffer(buft, 0) }; ++ ++ ggml_tensor * a = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 4, 4, 1); ++ ggml_tensor * mixes = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, 24, 1); ++ ggml_tensor * sc = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 3); ++ ggml_tensor * bs = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 24); ++ ggml_tensor * ik = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F16, 128, 16); ++ ggml_tensor * dq = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 128, 24, 1); ++ ggml_tensor * dw = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, 24, 1); ++ ggml_tensor * dm = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F16, 16, 1); ++ ++ ggml_tensor * op_sinkhorn = ggml_sinkhorn(ctx.get(), a, 20, 1e-6f); ++ ggml_tensor * op_hc_mix = ggml_hc_mix(ctx.get(), mixes, sc, bs, 4, 20, 1e-6f); ++ ggml_tensor * op_dsa_score = ggml_dsa_score(ctx.get(), ik, dq, dw, dm); ++ ++ a->buffer = buf.get(); ++ mixes->buffer = buf.get(); ++ sc->buffer = buf.get(); ++ bs->buffer = buf.get(); ++ ik->buffer = buf.get(); ++ dq->buffer = buf.get(); ++ dw->buffer = buf.get(); ++ dm->buffer = buf.get(); ++ ++ if (!ggml_backend_dev_supports_op(dev, op_sinkhorn) || ++ !ggml_backend_dev_supports_op(dev, op_hc_mix) || ++ !ggml_backend_dev_supports_op(dev, op_dsa_score)) { ++ LLAMA_LOG_WARN("%s: device %s does not support the fused hyper-connection ops - using the unfused op sequence\n", ++ __func__, ggml_backend_dev_name(dev)); ++ sinkhorn_fused_probe = 0; ++ break; ++ } ++ } ++ } ++ return sinkhorn_fused_probe != 0; ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_hc_sinkhorn(ggml_tensor * comb, int il) const { ++ GGML_UNUSED(il); ++ ++ if (use_fused_sinkhorn) { ++ // fused softmax + eps + alternating row/column normalization: one op instead ++ // of the ~8 nodes x 2 x n_iter the unfused form needs per [hc, hc] matrix ++ return ggml_sinkhorn(ctx0, comb, hparams.dsv4_hc_sinkhorn_iters, hparams.dsv4_hc_eps); ++ } ++ ++ // unfused fallback for backends without GGML_OP_SINKHORN (numerically equivalent) ++ comb = ggml_soft_max(ctx0, comb); ++ ++ ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1); ++ eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps); ++ comb = ggml_add(ctx0, comb, eps); ++ ++ auto norm_cols = [&]() { ++ ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); ++ ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst); ++ col_sum = ggml_add(ctx0, col_sum, eps); ++ col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3); ++ comb = ggml_div(ctx0, comb, col_sum); ++ }; ++ auto norm_rows = [&]() { ++ ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb); ++ row_sum = ggml_add(ctx0, row_sum, eps); ++ comb = ggml_div(ctx0, comb, row_sum); ++ }; ++ ++ norm_cols(); ++ for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) { ++ norm_rows(); ++ norm_cols(); ++ } ++ return comb; ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_hc_pre(ggml_tensor * x, ggml_tensor * xt, ++ ggml_tensor * hc_fn, ggml_tensor * hc_scale, ++ ggml_tensor * hc_base, ggml_tensor ** post, ggml_tensor ** comb, int il) const { ++ const int64_t hc = hparams.dsv4_hc_mult; ++ const int64_t hc_dim = hc * n_embd; ++ const int64_t nt = x->ne[2]; ++ ++ ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); ++ ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps); ++ ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); ++ ++ if (use_fused_sinkhorn) { ++ // fused affine + sigmoid gates + sinkhorn: one op instead of ~12 per call ++ ggml_tensor * mixed = ggml_hc_mix(ctx0, mixes, hc_scale, hc_base, ++ (int) hc, (int) hparams.dsv4_hc_sinkhorn_iters, hparams.dsv4_hc_eps); ++ ++ const size_t es = ggml_element_size(mixed); ++ ggml_tensor * pre = ggml_view_2d(ctx0, mixed, hc, nt, mixed->nb[1], 0); ++ *post = ggml_view_2d(ctx0, mixed, hc, nt, mixed->nb[1], hc*es); ++ *comb = ggml_view_3d(ctx0, mixed, hc, hc, nt, hc*es, mixed->nb[1], 2*hc*es); ++ ++ return build_hc_weighted_sum(x, xt, pre); ++ } ++ ++ ggml_tensor * scale_pre = opv2_view_1d(ctx0, hc_scale, 1, 0); ++ ggml_tensor * scale_post = opv2_view_1d(ctx0, hc_scale, 1, 1); ++ ggml_tensor * scale_comb = opv2_view_1d(ctx0, hc_scale, 1, 2); ++ ++ ggml_tensor * base_pre = opv2_view_1d(ctx0, hc_base, hc, 0); ++ ggml_tensor * base_post = opv2_view_1d(ctx0, hc_base, hc, hc); ++ ggml_tensor * base_comb = opv2_view_1d(ctx0, hc_base, hc * hc, 2 * hc); ++ ++ ggml_tensor * pre = opv2_view_2d(ctx0, mixes, hc, nt, 0); ++ pre = opv2_hc_affine(ctx0, pre, scale_pre, base_pre); ++ pre = ggml_sigmoid(ctx0, pre); ++ pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); ++ ++ *post = opv2_view_2d(ctx0, mixes, hc, nt, hc); ++ *post = opv2_hc_affine(ctx0, *post, scale_post, base_post); ++ *post = ggml_sigmoid(ctx0, *post); ++ *post = ggml_scale(ctx0, *post, 2.0f); ++ ++ *comb = opv2_view_2d(ctx0, mixes, hc * hc, nt, 2 * hc); ++ *comb = opv2_hc_affine(ctx0, *comb, scale_comb, base_comb); ++ *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt); ++ *comb = build_hc_sinkhorn(*comb, il); ++ ++ return build_hc_weighted_sum(x, xt, pre); ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_hc_post(ggml_tensor * x, ggml_tensor * residual, ++ ggml_tensor * residual_t, ggml_tensor * post, ggml_tensor * comb, int il) const { ++ GGML_UNUSED(il); ++ const int64_t hc = hparams.dsv4_hc_mult; ++ const int64_t nt = x->ne[1]; ++ ++ if (residual_t) { ++ // merged[e,dst,t] = sum_src residual_t[src,e,t] * comb[dst,src,t]: batched matmul. ++ // comb is laid out [dst, src, t]; mul_mat contracts dim0, so transpose the ++ // (tiny) mixing matrix to [src, dst, t] first ++ ggml_tensor * comb_t = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); ++ ggml_tensor * merged = ggml_mul_mat(ctx0, residual_t, comb_t); // [n_embd, hc, nt] ++ ++ // xpost[e,dst,t] = x[e,t] * post[dst,t]: rank-1 outer product per token ++ if (!ggml_is_contiguous(post)) { ++ post = ggml_cont(ctx0, post); // hc_mix output views are strided ++ } ++ ggml_tensor * xr = ggml_reshape_3d(ctx0, x, 1, n_embd, nt); ++ ggml_tensor * pr = ggml_reshape_3d(ctx0, post, 1, hc, nt); ++ ggml_tensor * xpost = ggml_mul_mat(ctx0, xr, pr); // [n_embd, hc, nt] ++ ++ return ggml_add(ctx0, merged, xpost); ++ } ++ ++ ggml_tensor * out = nullptr; ++ for (int64_t dst = 0; dst < hc; ++dst) { ++ ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst * post->nb[0]); ++ ggml_tensor * cur = ggml_mul(ctx0, x, post_dst); ++ for (int64_t src = 0; src < hc; ++src) { ++ ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src * residual->nb[1]); ++ ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], dst * comb->nb[0] + src * comb->nb[1]); ++ cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst)); ++ } ++ cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt); ++ out = out ? ggml_concat(ctx0, out, cur, 1) : cur; ++ } ++ return out; ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_hc_head(ggml_tensor * x, ggml_tensor * hc_fn, ++ ggml_tensor * hc_scale, ggml_tensor * hc_base) const { ++ const int64_t hc = hparams.dsv4_hc_mult; ++ const int64_t hc_dim = hc * n_embd; ++ const int64_t nt = x->ne[2]; ++ ++ ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); ++ ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps); ++ ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); ++ ++ // pre_only merge: plain sigmoid, no eps (the reference only adds hc_eps on ++ // the full pre/post/comb path, not in the pre_only head) ++ ggml_tensor * pre = opv2_hc_affine(ctx0, mixes, hc_scale, hc_base); ++ pre = ggml_sigmoid(ctx0, pre); ++ ++ ggml_tensor * xt = nt <= OPV2_HC_BATCHED_NT_MAX ++ ? ggml_cont(ctx0, ggml_permute(ctx0, x, 1, 0, 2, 3)) // [hc, n_embd, nt] ++ : nullptr; ++ return build_hc_weighted_sum(x, xt, pre); ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_mome_conv(ggml_tensor * x, ggml_tensor * conv_w, ++ ggml_tensor * conv_rs, ggml_tensor * conv_state, int64_t rs_head, ++ int64_t state_off, int64_t n_seqs, int64_t n_seq_tokens) const { ++ // x: [dim, n_tokens]; conv_w: [k, dim] (bf16 in gguf, cast to f32 for ssm_conv). ++ // Causal depthwise conv with an identity residual. The previous (k-1) input ++ // columns are carried across ubatches via the recurrent state. ++ const int64_t dim = x->ne[0]; ++ const int64_t nt = x->ne[1]; ++ const int64_t k = conv_w->ne[0]; ++ const int64_t d_conv = k - 1; // state width per channel ++ const int64_t esz = ggml_element_size(conv_state); ++ ++ ggml_tensor * w = conv_w->type == GGML_TYPE_F32 ? conv_w : ggml_cast(ctx0, conv_w, GGML_TYPE_F32); ++ ++ // previous-token state for this conv: strided view of the gathered recurrent state ++ ggml_tensor * state = ggml_view_3d(ctx0, conv_rs, d_conv, dim, n_seqs, ++ d_conv * esz, conv_rs->nb[1], state_off * esz); // [k-1, dim, n_seqs] ++ ++ // input as [n_seq_tokens, dim, n_seqs]; concat handles the non-contiguous permute ++ ggml_tensor * xt = ggml_reshape_3d(ctx0, x, dim, n_seq_tokens, n_seqs); ++ xt = ggml_permute(ctx0, xt, 1, 0, 2, 3); // [n_seq_tokens, dim, n_seqs] ++ ++ ggml_tensor * sx = ggml_concat(ctx0, state, xt, 0); // [k-1+n_seq_tokens, dim, n_seqs] ++ ++ // store the last (k-1) input columns as the next state ++ ggml_tensor * new_state = ggml_view_3d(ctx0, sx, d_conv, dim, n_seqs, sx->nb[1], sx->nb[2], ++ n_seq_tokens * sx->nb[0]); ++ ggml_tensor * dst = ggml_view_2d(ctx0, conv_state, d_conv * dim, n_seqs, conv_state->nb[1], ++ rs_head * conv_state->nb[1] + state_off * esz); ++ ggml_build_forward_expand(gf, ggml_cpy(ctx0, ++ ggml_reshape_2d(ctx0, ggml_cont(ctx0, new_state), d_conv * dim, n_seqs), dst)); ++ ++ ggml_tensor * conv = ggml_ssm_conv(ctx0, sx, w); // [dim, n_seq_tokens, n_seqs] ++ conv = ggml_reshape_2d(ctx0, conv, dim, nt); ++ ++ return ggml_add(ctx0, x, conv); // identity residual ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_attention(const llama_model & model, llm_graph_input_mem_hybrid_iswa * inp, ++ ggml_tensor * cur, ggml_tensor * inp_pos, float kq_scale, int il) const { ++ const auto & layer = model.layers[il]; ++ auto * inp_attn = inp->get_attn(); ++ const bool is_swa = hparams.is_swa(il); ++ ++ const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); ++ const int64_t n_embd_head_qk_rope = hparams.n_rot(); ++ const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope; ++ const int64_t kv_lora_rank = hparams.n_lora_kv; ++ const int64_t nt = cur->ne[1]; ++ ++ // recurrent MoME conv state (previous k-1 input columns per conv, per sequence) ++ const int64_t n_seqs = ubatch.n_seqs; ++ const int64_t n_seq_tokens = ubatch.n_seq_tokens; ++ const int64_t d_conv = hparams.openpangu_conv_k - 1; ++ const int64_t off_qa = 0; ++ const int64_t off_kv = d_conv * hparams.n_lora_q; ++ const int64_t off_o = d_conv * (hparams.n_lora_q + hparams.n_lora_kv); ++ ++ const auto * rmctx = static_cast(mctx)->get_recr(); ++ ggml_tensor * conv_state = rmctx->get_r_l(il); ++ const int64_t rs_head = rmctx->get_head(); ++ ggml_tensor * conv_rs = build_rs(inp->get_recr(), conv_state, hparams.n_embd_r(), n_seqs); ++ ++ // q down -> conv -> norm -> up ++ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur); ++ q = build_mome_conv(q, layer.attn_q_a_conv, conv_rs, conv_state, rs_head, off_qa, n_seqs, n_seq_tokens); ++ ggml_tensor * q_lora = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); ++ q = ggml_mul_mat(ctx0, layer.wq_b, q_lora); ++ ++ ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, nt, ++ ggml_row_size(q->type, n_embd_head_k), ggml_row_size(q->type, n_embd_head_k) * n_head, 0); ++ ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, nt, ++ ggml_row_size(q->type, n_embd_head_k), ggml_row_size(q->type, n_embd_head_k) * n_head, ++ ggml_row_size(q->type, n_embd_head_qk_nope)); ++ ++ // kv down: split compressed (kv_lora) and rope (k_pe); conv applies to compressed only ++ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); ++ ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, nt, ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); ++ ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, nt, ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); ++ ++ kv_cmpr = ggml_cont(ctx0, kv_cmpr); ++ kv_cmpr = build_mome_conv(kv_cmpr, layer.attn_kv_a_conv, conv_rs, conv_state, rs_head, off_kv, n_seqs, n_seq_tokens); ++ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); ++ ++ q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_embd_head_qk_rope, GGML_ROPE_TYPE_NEOX, ++ n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); ++ k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_embd_head_qk_rope, GGML_ROPE_TYPE_NEOX, ++ n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); ++ ++ // absorb q_nope into the compressed space via wk_b ++ q_nope = ggml_cont(ctx0, ggml_permute(ctx0, q_nope, 0, 2, 1, 3)); // [nope, nt, n_head] ++ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope); ++ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3); // [kv_lora, n_head, nt] ++ ++ ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0); // [kv_lora+rope, n_head, nt] ++ ++ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, nt); ++ ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0); // [kv_lora+rope, 1, nt] ++ ++ // DSA lightning indexer: the per-token indexer key rides in the tail of the ++ // cached K row; layers without an indexer (SWA) zero-pad the tail instead ++ const int64_t idx_d = hparams.indexer_head_size; ++ const int64_t idx_h = hparams.indexer_n_head; ++ const bool is_dsa = idx_d > 0 && hparams.indexer_top_k > 0 && !is_swa && ++ layer.indexer_attn_q_b && layer.indexer_attn_k && layer.indexer_proj; ++ ++ ggml_tensor * indexer_q = nullptr; ++ if (is_dsa) { ++ // query heads from the post-conv, post-norm q_lora (the wq_b input); ++ // rope on the leading n_rot dims of each head, as in the main attention ++ indexer_q = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, q_lora); // [idx_h*idx_d, nt] ++ ggml_tensor * iq_pe = ggml_view_3d(ctx0, indexer_q, n_embd_head_qk_rope, idx_h, nt, ++ ggml_row_size(indexer_q->type, idx_d), ggml_row_size(indexer_q->type, idx_d) * idx_h, 0); ++ ggml_tensor * iq_nope = ggml_view_3d(ctx0, indexer_q, idx_d - n_embd_head_qk_rope, idx_h, nt, ++ ggml_row_size(indexer_q->type, idx_d), ggml_row_size(indexer_q->type, idx_d) * idx_h, ++ ggml_row_size(indexer_q->type, n_embd_head_qk_rope)); ++ iq_pe = ggml_rope_ext(ctx0, iq_pe, inp_pos, nullptr, n_embd_head_qk_rope, GGML_ROPE_TYPE_NEOX, ++ n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); ++ indexer_q = ggml_concat(ctx0, iq_pe, iq_nope, 0); // [idx_d, idx_h, nt] ++ ++ // single shared key from the attention input: RMS norm, same rope split ++ ggml_tensor * ik = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur); // [idx_d, nt] ++ ik = build_norm(ik, layer.indexer_k_norm, nullptr, LLM_NORM_RMS, il); ++ ggml_tensor * ik_pe = ggml_view_3d(ctx0, ik, n_embd_head_qk_rope, 1, nt, ++ ggml_row_size(ik->type, idx_d), ggml_row_size(ik->type, idx_d), 0); ++ ggml_tensor * ik_nope = ggml_view_3d(ctx0, ik, idx_d - n_embd_head_qk_rope, 1, nt, ++ ggml_row_size(ik->type, idx_d), ggml_row_size(ik->type, idx_d), ++ ggml_row_size(ik->type, n_embd_head_qk_rope)); ++ ik_pe = ggml_rope_ext(ctx0, ik_pe, inp_pos, nullptr, n_embd_head_qk_rope, GGML_ROPE_TYPE_NEOX, ++ n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); ++ ik = ggml_concat(ctx0, ik_pe, ik_nope, 0); // [idx_d, 1, nt] ++ ++ Kcur = ggml_concat(ctx0, Kcur, ik, 0); // [kv_lora+rope+idx_d, 1, nt] ++ } else if (idx_d > 0) { ++ Kcur = ggml_pad(ctx0, Kcur, idx_d, 0, 0, 0); ++ } ++ ++ // store compressed K to the right iSWA sub-cache (full for DSA/global layers, ++ // windowed for SWA layers). MLA: V is a view of K, decompressed by wv_b. ++ ggml_build_forward_expand(gf, Qcur); ++ ggml_build_forward_expand(gf, Kcur); ++ const auto * attn_mctx = is_swa ? inp_attn->mctx->get_swa() : inp_attn->mctx->get_base(); ++ ggml_tensor * k_idxs = is_swa ? inp_attn->get_k_idxs_swa() : inp_attn->get_k_idxs(); ++ ggml_build_forward_expand(gf, attn_mctx->cpy_k(ctx0, Kcur, k_idxs, il)); ++ ++ const int64_t n_pfx = hparams.n_k_sink_prefix; ++ ++ ggml_tensor * k_row = attn_mctx->get_k(ctx0, il); // [kv_lora+rope(+idx_d), 1, n_pfx+n_kv] ++ ++ // cache-resident param sinks: refresh the reserved prefix rows (cheap, static values) ++ build_sink_write(layer, k_row, il); ++ ++ ggml_tensor * k = idx_d == 0 ? k_row ++ : ggml_view_4d(ctx0, k_row, kv_lora_rank + n_embd_head_qk_rope, k_row->ne[1], k_row->ne[2], k_row->ne[3], ++ k_row->nb[1], k_row->nb[2], k_row->nb[3], 0); // MLA part: [kv_lora+rope, 1, n_pfx+n_kv] ++ ++ ggml_tensor * v = ggml_view_4d(ctx0, k_row, kv_lora_rank, k_row->ne[1], k_row->ne[2], k_row->ne[3], ++ k_row->nb[1], k_row->nb[2], k_row->nb[3], 0); ++ ++ // hoisted combined mask [n_pfx | n_kv]; DSA layers refine the cell part below ++ ggml_tensor * kq_mask_all = is_swa ? kq_mask_sinked_swa : kq_mask_sinked; ++ ggml_tensor * kq_mask = is_swa ? inp_attn->get_kq_mask_swa() : inp_attn->get_kq_mask(); ++ ++ // while the whole cache fits in the top-k budget the selection is a no-op: ++ // skip the scoring (the indexer keys are still written to the cache above) ++ if (is_dsa && k_row->ne[2] - n_pfx > (int64_t) hparams.indexer_top_k) { ++ // score the cached indexer keys (cell rows only; the param sinks are not ++ // scored -- they stay always-visible through the mask prefix): ++ // score[t,s] = sum_h w[t,h] * relu(q[t,h,:] . k[s,:]) (no scaling, per the reference) ++ ggml_tensor * ik_cache = ggml_view_4d(ctx0, k_row, idx_d, k_row->ne[1], k_row->ne[2] - n_pfx, k_row->ne[3], ++ k_row->nb[1], k_row->nb[2], k_row->nb[3], ++ n_pfx*k_row->nb[2] + ggml_row_size(k_row->type, kv_lora_rank + n_embd_head_qk_rope)); // [idx_d, 1, n_kv] ++ ++ ggml_tensor * iw = ggml_mul_mat(ctx0, layer.indexer_proj, cur); // [idx_h, nt] ++ ++ ggml_tensor * iscore; ++ if (use_fused_sinkhorn && k_row->ne[3] == 1) { ++ // one fused pass over the cached indexer keys (incl. relu, head-weighting ++ // and the additive mask); the CUDA backend has a warp-per-position decode ++ // kernel and a tiled prefill kernel, so no [n_kv, nh, nt] intermediate is ++ // ever materialized. multi-stream caches fall back to the matmul pipeline. ++ ggml_tensor * ikv = ggml_view_2d(ctx0, k_row, idx_d, k_row->ne[2] - n_pfx, k_row->nb[2], ++ n_pfx*k_row->nb[2] + ggml_row_size(k_row->type, kv_lora_rank + n_embd_head_qk_rope)); ++ iscore = ggml_dsa_score(ctx0, ikv, indexer_q, iw, kq_mask); // [n_kv, nt] ++ } else { ++ // prefill: tiled matmul pipeline ++ // split the batch into streams if needed ++ const int64_t n_stream = ik_cache->ne[3]; ++ indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, ++ indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0); ++ iw = ggml_view_4d(ctx0, iw, iw->ne[0], iw->ne[1]/n_stream, iw->ne[2], n_stream, ++ iw->nb[1], iw->nb[2]/n_stream, iw->nb[3]/n_stream, 0); ++ ++ indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); // [idx_d, nt, idx_h] ++ ik_cache = ggml_permute(ctx0, ik_cache, 0, 2, 1, 3); // [idx_d, n_kv, 1] ++ ++ ggml_tensor * ikq = ggml_mul_mat(ctx0, ik_cache, indexer_q); // [n_kv, nt, idx_h] ++ ikq = ggml_relu(ctx0, ikq); // contiguous, no cont needed ++ ikq = ggml_cont(ctx0, ggml_permute(ctx0, ikq, 1, 2, 0, 3)); // [idx_h, n_kv, nt] ++ ++ // weighted sum over the heads as a batched matvec ++ ggml_tensor * iwv = ggml_view_4d(ctx0, iw, iw->ne[0], 1, iw->ne[1], iw->ne[3], ++ iw->nb[2], iw->nb[1], iw->nb[3], 0); // [idx_h, 1, nt] ++ iscore = ggml_mul_mat(ctx0, ikq, iwv); // [n_kv, 1, nt] ++ iscore = ggml_reshape_4d(ctx0, iscore, iscore->ne[0], iscore->ne[2], 1, iscore->ne[3]); // [n_kv, nt, 1] ++ ++ // mask before top-k so invalid/future cells cannot claim slots ++ iscore = ggml_add(ctx0, iscore, ++ kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32)); ++ } ++ ++ const int64_t n_top_k = std::min(iscore->ne[0], hparams.indexer_top_k); ++ // GGML_OP_TOP_K output is already contiguous (the cont dated from the argsort+view form) ++ ggml_tensor * top_k = ggml_top_k(ctx0, iscore, n_top_k); ++ ++ if (nt == 1) { ++ // gather-based decode attention: physically gather the selected rows and ++ // attend over [sinks | top-k] only -- attention cost stays constant with ++ // context instead of reading (and masking) the entire cache ++ ggml_tensor * ids = ggml_reshape_1d(ctx0, top_k, n_top_k); ++ ++ ggml_tensor * cells = ggml_view_2d(ctx0, k_row, k_row->ne[0], k_row->ne[2] - n_pfx, ++ k_row->nb[2], n_pfx*k_row->nb[2]); // [row_w, n_kv] ++ ggml_tensor * gk = ggml_get_rows(ctx0, cells, ids); // [row_w, n_top_k] F32 ++ gk = ggml_cast(ctx0, gk, k_row->type); ++ ++ ggml_tensor * prefix = ggml_view_2d(ctx0, k_row, k_row->ne[0], n_pfx, k_row->nb[2], 0); ++ ggml_tensor * kg = ggml_concat(ctx0, prefix, gk, 1); // [row_w, n_pfx+n_top_k] ++ ++ const int64_t n_rows = n_pfx + n_top_k; ++ k = ggml_view_4d(ctx0, kg, kv_lora_rank + n_embd_head_qk_rope, 1, n_rows, 1, ++ kg->nb[1], kg->nb[1], kg->nb[1]*n_rows, 0); ++ v = ggml_view_4d(ctx0, kg, kv_lora_rank, 1, n_rows, 1, ++ kg->nb[1], kg->nb[1], kg->nb[1]*n_rows, 0); ++ ++ // gather the mask entries of the selected cells (keeps padding/other-seq ++ // cells that sneak into the top-k masked out) ++ ggml_tensor * m_r = ggml_reshape_2d(ctx0, kq_mask, 1, kq_mask->ne[0]); ++ ggml_tensor * gm = ggml_get_rows(ctx0, m_r, ids); // [1, n_top_k] F32 ++ gm = ggml_reshape_2d(ctx0, gm, n_top_k, 1); ++ if (gm->type != sink_mask_blk->type) { ++ gm = ggml_cast(ctx0, gm, sink_mask_blk->type); ++ } ++ kq_mask_all = ggml_concat(ctx0, sink_mask_blk, gm, 0); // [n_pfx+n_top_k, 1] ++ } else { ++ // prefill: every query has its own top-k set, so keep the mask-based form -- ++ // unmask the selected cells on an all-masked copy, then AND with the causal ++ // mask (same construction as the DSA build_attn overload in llama-graph.cpp) ++ ggml_tensor * mask_inf = ggml_fill(ctx0, kq_mask, -INFINITY); ++ mask_inf = ggml_view_4d(ctx0, mask_inf, 1, mask_inf->ne[0], mask_inf->ne[1], mask_inf->ne[3], ++ mask_inf->nb[0], mask_inf->nb[1], mask_inf->nb[2], 0); ++ ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, ++ top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); ++ ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); ++ zeros = ggml_fill(ctx0, zeros, 0.0f); ++ ggml_tensor * unmasked = ggml_set_rows(ctx0, mask_inf, zeros, top_k_3d); ++ unmasked = ggml_view_4d(ctx0, unmasked, unmasked->ne[1], unmasked->ne[2], 1, unmasked->ne[3], ++ unmasked->nb[2], unmasked->nb[3], unmasked->nb[3], 0); ++ kq_mask = ggml_add(ctx0, unmasked, kq_mask); ++ ++ // per-DSA-layer combined mask (the top-k selection differs per layer) ++ kq_mask_all = ggml_concat(ctx0, sink_mask_blk, kq_mask, 0); ++ } ++ } ++ ++ ggml_tensor * attn_out = build_attn_mha(Qcur, k, v, nullptr, kq_mask_all, nullptr, layer.wv_b, 0, kq_scale, il); ++ ++ attn_out = build_mome_conv(attn_out, layer.attn_o_conv, conv_rs, conv_state, rs_head, off_o, n_seqs, n_seq_tokens); ++ attn_out = ggml_mul_mat(ctx0, layer.wo, attn_out); // o_proj ++ return attn_out; ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_attention_mtp(const llama_model & model, llm_graph_input_attn_kv_iswa * inp_attn, ++ ggml_tensor * cur, ggml_tensor * inp_pos, float kq_scale, int il) const { ++ const auto & layer = model.layers[il]; ++ // in the combined GGUF the MTP layers are converted as full attention (the single ++ // global n_swa cannot express the reference's 2048 next to the trunk's 512); a ++ // split MTP-only GGUF carries its own n_swa = 2048 and runs SWA here, matching ++ // the reference -- see conversion/pangu.py and conversion/split_pangu_mtp.py ++ const bool is_swa = hparams.is_swa(il); ++ ++ const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); ++ const int64_t n_embd_head_qk_rope = hparams.n_rot(); ++ const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope; ++ const int64_t kv_lora_rank = hparams.n_lora_kv; ++ const int64_t nt = cur->ne[1]; ++ ++ // q down -> norm -> up (no MoME conv on the MTP head) ++ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur); ++ q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); ++ q = ggml_mul_mat(ctx0, layer.wq_b, q); ++ ++ ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, nt, ++ ggml_row_size(q->type, n_embd_head_k), ggml_row_size(q->type, n_embd_head_k) * n_head, 0); ++ ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, nt, ++ ggml_row_size(q->type, n_embd_head_k), ggml_row_size(q->type, n_embd_head_k) * n_head, ++ ggml_row_size(q->type, n_embd_head_qk_nope)); ++ ++ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); ++ ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, nt, ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); ++ ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, nt, ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); ++ ++ kv_cmpr = ggml_cont(ctx0, kv_cmpr); ++ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); ++ ++ q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_embd_head_qk_rope, GGML_ROPE_TYPE_NEOX, ++ n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); ++ k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_embd_head_qk_rope, GGML_ROPE_TYPE_NEOX, ++ n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); ++ ++ q_nope = ggml_cont(ctx0, ggml_permute(ctx0, q_nope, 0, 2, 1, 3)); ++ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope); ++ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3); ++ ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0); ++ ++ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, nt); ++ ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0); ++ ++ // MTP layers have no indexer -- zero-pad the DSA tail of the cache row ++ const int64_t idx_d = hparams.indexer_head_size; ++ if (idx_d > 0) { ++ Kcur = ggml_pad(ctx0, Kcur, idx_d, 0, 0, 0); ++ } ++ ++ ggml_build_forward_expand(gf, Qcur); ++ ggml_build_forward_expand(gf, Kcur); ++ const auto * attn_mctx = is_swa ? inp_attn->mctx->get_swa() : inp_attn->mctx->get_base(); ++ ggml_tensor * k_idxs = is_swa ? inp_attn->get_k_idxs_swa() : inp_attn->get_k_idxs(); ++ ggml_build_forward_expand(gf, attn_mctx->cpy_k(ctx0, Kcur, k_idxs, il)); ++ ++ ggml_tensor * k_row = attn_mctx->get_k(ctx0, il); // [row_w, 1, n_pfx+n_kv] ++ ++ // cache-resident param sinks: refresh the reserved prefix rows ++ build_sink_write(layer, k_row, il); ++ ++ ggml_tensor * k = idx_d == 0 ? k_row ++ : ggml_view_4d(ctx0, k_row, kv_lora_rank + n_embd_head_qk_rope, k_row->ne[1], k_row->ne[2], k_row->ne[3], ++ k_row->nb[1], k_row->nb[2], k_row->nb[3], 0); ++ ++ ggml_tensor * v = ggml_view_4d(ctx0, k_row, kv_lora_rank, k_row->ne[1], k_row->ne[2], k_row->ne[3], ++ k_row->nb[1], k_row->nb[2], k_row->nb[3], 0); ++ ++ // hoisted combined mask ++ ggml_tensor * kq_mask_all = is_swa ? kq_mask_sinked_swa : kq_mask_sinked; ++ ++ ggml_tensor * attn_out = build_attn_mha(Qcur, k, v, nullptr, kq_mask_all, nullptr, layer.wv_b, 0, kq_scale, il); ++ attn_out = ggml_mul_mat(ctx0, layer.wo, attn_out); // o_proj (no MoME conv) ++ return attn_out; ++} ++ ++ggml_tensor * llama_model_openpangu_v2::graph_base::build_moe_block(const llama_model & model, ggml_tensor * cur, int il) const { ++ if (il < (int) hparams.n_layer_dense_lead) { ++ return build_ffn(cur, ++ model.layers[il].ffn_up, nullptr, nullptr, ++ model.layers[il].ffn_gate, nullptr, nullptr, ++ model.layers[il].ffn_down, nullptr, nullptr, ++ nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); ++ } ++ ggml_tensor * moe_out = build_moe_ffn(cur, ++ model.layers[il].ffn_gate_inp, ++ model.layers[il].ffn_up_exps, ++ model.layers[il].ffn_gate_exps, ++ model.layers[il].ffn_down_exps, ++ model.layers[il].ffn_exp_probs_b, ++ n_expert, n_expert_used, ++ LLM_FFN_SILU, hparams.expert_weights_norm, ++ hparams.expert_weights_scale, ++ (llama_expert_gating_func_type) hparams.expert_gating_func, ++ il); ++ ggml_tensor * ffn_shexp = build_ffn(cur, ++ model.layers[il].ffn_up_shexp, nullptr, nullptr, ++ model.layers[il].ffn_gate_shexp, nullptr, nullptr, ++ model.layers[il].ffn_down_shexp, nullptr, nullptr, ++ nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); ++ return ggml_add(ctx0, moe_out, ffn_shexp); ++} ++ ++std::unique_ptr llama_model_openpangu_v2::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); ++} ++ ++llama_model_openpangu_v2::graph::graph(const llama_model & model, const llm_graph_params & params) : ++ graph_base(params, static_cast(model).sinkhorn_fused()) { ++ GGML_ASSERT(n_layer > 0 && "openpangu-v2: no trunk layers -- a split MTP-only GGUF can only be used as --model-draft with --mtp"); ++ ++ const int64_t hc = hparams.dsv4_hc_mult; ++ const float kq_scale = 1.0f / sqrtf(float(hparams.n_embd_head_k_mla())); ++ ++ ggml_tensor * cur; ++ ggml_tensor * inp = build_inp_embd(model.tok_embd); ++ ggml_tensor * inp_pos = build_inp_pos(); ++ ggml_tensor * inp_out_ids = build_inp_out_ids(); ++ ++ auto * inp_hybrid = build_inp_mem_hybrid_iswa(); ++ ++ // hoisted sink masks: one [n_pfx | n_kv] concat per variant per graph ++ { ++ ggml_tensor * mask_base = inp_hybrid->get_attn()->get_kq_mask(); ++ ggml_tensor * mask_swa = inp_hybrid->get_attn()->get_kq_mask_swa(); ++ sink_mask_blk = build_sink_mask_blk(mask_base); ++ kq_mask_sinked = ggml_concat(ctx0, sink_mask_blk, mask_base, 0); ++ kq_mask_sinked_swa = ggml_concat(ctx0, build_sink_mask_blk(mask_swa), mask_swa, 0); ++ } ++ ++ ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens); ++ inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); ++ ++ // transposed stream state [hc, n_embd, nt] for the batched-matmul mHC form, ++ // shared by the pre- and post-merge of a block; null selects the loop form ++ auto make_xt = [&](ggml_tensor * s) -> ggml_tensor * { ++ return s->ne[2] <= OPV2_HC_BATCHED_NT_MAX ++ ? ggml_cont(ctx0, ggml_permute(ctx0, s, 1, 0, 2, 3)) ++ : nullptr; ++ }; ++ ++ for (int il = 0; il < n_layer; ++il) { ++ ggml_tensor * post = nullptr; ++ ggml_tensor * comb = nullptr; ++ ++ // attention block ++ ggml_tensor * residual = inpL; ++ ggml_tensor * residual_t = make_xt(inpL); ++ cur = build_hc_pre(inpL, residual_t, model.layers[il].hc_attn_fn, model.layers[il].hc_attn_scale, ++ model.layers[il].hc_attn_base, &post, &comb, il); ++ cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); ++ cur = build_attention(model, inp_hybrid, cur, inp_pos, kq_scale, il); ++ cur = build_norm(cur, model.layers[il].attn_post_norm, nullptr, LLM_NORM_RMS, il); ++ inpL = build_hc_post(cur, residual, residual_t, post, comb, il); ++ ++ // the last layer's FFN only needs the requested output rows -- gather early, ++ // unless the unmasked MTP extraction needs the full-length hidden state ++ if (il == n_layer - 1 && inp_out_ids && ++ (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked)) { ++ 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); ++ } ++ ++ // feed-forward block ++ residual = inpL; ++ residual_t = make_xt(inpL); ++ cur = build_hc_pre(inpL, residual_t, model.layers[il].hc_ffn_fn, model.layers[il].hc_ffn_scale, ++ model.layers[il].hc_ffn_base, &post, &comb, il); ++ cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il); ++ ++ cur = build_moe_block(model, cur, il); ++ ++ cur = build_norm(cur, model.layers[il].ffn_post_norm, nullptr, LLM_NORM_RMS, il); ++ inpL = build_hc_post(cur, residual, residual_t, post, comb, il); ++ ++ // block post norm on the stream-concatenated residual (subset of layers) ++ if (model.layers[il].block_post_norm) { ++ const int64_t nt_l = inpL->ne[2]; ++ ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd * hc, nt_l); ++ flat = build_norm(flat, model.layers[il].block_post_norm, nullptr, LLM_NORM_RMS, il); ++ inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, nt_l); ++ } ++ ++ inpL = build_cvec(inpL, il); ++ } ++ ++ // Collapse the mHC streams. In the unmasked-MTP case inpL still holds all ++ // n_tokens rows so t_h_nextn covers every position; otherwise the last-layer ++ // gather above already reduced it to the output rows. ++ cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base); ++ ++ // trunk hidden state (pre-output-norm), consumed by the MTP head for speculative drafting ++ cb(cur, "h_nextn", -1); ++ res->t_h_nextn = cur; ++ ++ if (cparams.embeddings_nextn && !cparams.embeddings_nextn_masked && inp_out_ids) { ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ } ++ ++ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); ++ res->t_embd = cur; ++ ++ cur = ggml_mul_mat(ctx0, model.output, cur); ++ res->t_logits = cur; ++ ++ ggml_build_forward_expand(gf, cur); ++} ++ ++llama_model_openpangu_v2::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : ++ graph_base(params, static_cast(model).sinkhorn_fused()) { ++ GGML_ASSERT(hparams.n_layer_nextn > 0 && "openpangu-v2 MTP requires n_layer_nextn > 0"); ++ ++ const int il = hparams.n_layer() + cparams.nextn_layer_offset; ++ GGML_ASSERT(cparams.nextn_layer_offset >= 0 && ++ cparams.nextn_layer_offset < (int) hparams.n_layer_nextn); ++ const auto & layer = model.layers[il]; ++ GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm); ++ ++ const float kq_scale = 1.0f / sqrtf(float(hparams.n_embd_head_k_mla())); ++ ++ // inputs: the next token id + the previous hidden state (mtp_h_input) ++ auto inp = std::make_unique(hparams.n_embd); ++ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); ++ ggml_set_input(inp->tokens); ++ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); ++ ggml_set_input(inp->embd); ++ ggml_set_name(inp->embd, "mtp_h_input"); ++ ++ ggml_tensor * h_input = inp->embd; ++ ggml_tensor * embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; ++ ggml_tensor * tok_embd = ggml_get_rows(ctx0, embd_w, inp->tokens); ++ res->add_input(std::move(inp)); ++ ++ ggml_tensor * inp_pos = build_inp_pos(); ++ auto * inp_attn = build_attn_inp_kv_iswa(); ++ ++ // hoisted sink masks (the swa variant is live for split MTP-only models, where ++ // the MTP layers run their reference SWA-2048 window) ++ { ++ ggml_tensor * mask_base = inp_attn->get_kq_mask(); ++ ggml_tensor * mask_swa = inp_attn->get_kq_mask_swa(); ++ sink_mask_blk = build_sink_mask_blk(mask_base); ++ kq_mask_sinked = ggml_concat(ctx0, sink_mask_blk, mask_base, 0); ++ kq_mask_sinked_swa = ggml_concat(ctx0, build_sink_mask_blk(mask_swa), mask_swa, 0); ++ } ++ ++ // eh_proj( concat[ enorm(embed(next_tok)), hnorm(prev_hidden) ] ) ++ ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); ++ ggml_tensor * h_norm = build_norm(h_input, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); ++ ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0)); ++ ++ // plain (non-mHC) decoder layer: sandwich norms + MLA(+sinks, no conv) + MoE ++ ggml_tensor * res_sa = cur; ++ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); ++ cur = build_attention_mtp(model, inp_attn, cur, inp_pos, kq_scale, il); ++ cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il); ++ cur = ggml_add(ctx0, cur, res_sa); ++ ++ ggml_tensor * res_ffn = cur; ++ cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); ++ cur = build_moe_block(model, cur, il); ++ cur = build_norm(cur, layer.ffn_post_norm, nullptr, LLM_NORM_RMS, il); ++ cur = ggml_add(ctx0, cur, res_ffn); ++ ++ // hidden state to seed the next chained MTP head: keep all n_tokens rows, ++ // the unmasked nextn extraction reads them densely by token position ++ cb(cur, "h_nextn", -1); ++ res->t_h_nextn = cur; ++ ++ ggml_tensor * inp_out_ids = build_inp_out_ids(); ++ if (inp_out_ids) { ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ } ++ ++ ggml_tensor * head_norm = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm; ++ ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; ++ cur = build_norm(cur, head_norm, nullptr, LLM_NORM_RMS, -1); ++ cur = build_lora_mm(head_w, cur); ++ res->t_logits = cur; ++ ++ ggml_build_forward_expand(gf, cur); ++} +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index ba7f7eb4e..05a0c7d27 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -8821,6 +8821,91 @@ struct test_fill : public test_case { + } + }; + ++// GGML_OP_SINKHORN ++struct test_sinkhorn : public test_case { ++ const std::array ne; ++ const int n_iter; ++ const float eps; ++ ++ std::string vars() override { return VARS_TO_STR3(ne, n_iter, eps); } ++ ++ test_sinkhorn(std::array ne = { 4, 4, 64, 2 }, int n_iter = 20, float eps = 1e-6f) ++ : ne(ne), n_iter(n_iter), eps(eps) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], ne[3]); ++ ggml_set_name(a, "a"); ++ ++ ggml_tensor * out = ggml_sinkhorn(ctx, a, n_iter, eps); ++ ++ ggml_set_name(out, "out"); ++ ++ return out; ++ } ++}; ++ ++// GGML_OP_HC_MIX ++struct test_hc_mix : public test_case { ++ const int hc; ++ const int64_t nt; ++ const int n_iter; ++ const float eps; ++ ++ std::string vars() override { return VARS_TO_STR4(hc, nt, n_iter, eps); } ++ ++ test_hc_mix(int hc = 4, int64_t nt = 37, int n_iter = 20, float eps = 1e-6f) ++ : hc(hc), nt(nt), n_iter(n_iter), eps(eps) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ const int64_t n0 = 2*hc + hc*hc; ++ ++ ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n0, nt); ++ ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 3); ++ ggml_tensor * base = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n0); ++ ggml_set_name(a, "a"); ++ ggml_set_name(scale, "scale"); ++ ggml_set_name(base, "base"); ++ ++ ggml_tensor * out = ggml_hc_mix(ctx, a, scale, base, hc, n_iter, eps); ++ ++ ggml_set_name(out, "out"); ++ ++ return out; ++ } ++}; ++ ++// GGML_OP_DSA_SCORE ++struct test_dsa_score : public test_case { ++ const int64_t nd; ++ const int64_t nkv; ++ const int64_t nh; ++ const int64_t nt; ++ const ggml_type type_ik; ++ ++ std::string vars() override { return VARS_TO_STR5(nd, nkv, nh, nt, type_ik); } ++ ++ test_dsa_score(int64_t nd = 128, int64_t nkv = 4096, int64_t nh = 24, int64_t nt = 1, ++ ggml_type type_ik = GGML_TYPE_F16) ++ : nd(nd), nkv(nkv), nh(nh), nt(nt), type_ik(type_ik) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * ik = ggml_new_tensor_2d(ctx, type_ik, nd, nkv); ++ ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, nd, nh, nt); ++ ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nh, nt); ++ ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nkv, nt); ++ ggml_set_name(ik, "ik"); ++ ggml_set_name(q, "q"); ++ ggml_set_name(w, "w"); ++ ggml_set_name(mask, "mask"); ++ ++ ggml_tensor * out = ggml_dsa_score(ctx, ik, q, w, mask); ++ ++ ggml_set_name(out, "out"); ++ ++ return out; ++ } ++}; ++ + // GGML_OP_SOLVE_TRI + struct test_solve_tri : public test_case { + const ggml_type type; +@@ -11425,6 +11510,7 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_top_k_inf({248320, 1, 1, 1}, k)); + } + ++ test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {102400, 8, 1, 1}, 2048)); // DSA indexer selection (CUDA radix select) + for (int i = 0; i < 20; ++i) { + for (int k : {1, 2, 3, 7, 15, 100, 500, 1023, 9999}) { + if (k <= 1<> make_test_cases_eval() { + test_cases.emplace_back(new test_fill(0.0f)); + test_cases.emplace_back(new test_fill(2.0f, GGML_TYPE_F32, { 303, 207, 11, 3 })); + test_cases.emplace_back(new test_fill(-152.0f, GGML_TYPE_F32, { 800, 600, 4, 4 })); ++ ++ test_cases.emplace_back(new test_sinkhorn()); ++ test_cases.emplace_back(new test_sinkhorn({ 4, 4, 512, 1 }, 20, 1e-6f)); ++ test_cases.emplace_back(new test_sinkhorn({ 3, 5, 33, 2 }, 7, 1e-5f)); ++ test_cases.emplace_back(new test_sinkhorn({ 8, 8, 1, 1 }, 1, 1e-6f)); ++ ++ test_cases.emplace_back(new test_hc_mix()); ++ test_cases.emplace_back(new test_hc_mix(4, 512, 20, 1e-6f)); ++ test_cases.emplace_back(new test_hc_mix(2, 1, 5, 1e-5f)); ++ ++ test_cases.emplace_back(new test_dsa_score()); ++ test_cases.emplace_back(new test_dsa_score(128, 2560, 24, 4, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_dsa_score(64, 1000, 8, 2, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_dsa_score(128, 102400, 24, 1, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_dsa_score(128, 4096, 24, 512, GGML_TYPE_F16)); // CUDA tiled prefill kernel ++ test_cases.emplace_back(new test_dsa_score(128, 1000, 24, 100, GGML_TYPE_F16)); // tiled, ragged tiles ++ test_cases.emplace_back(new test_dsa_score(64, 999, 8, 37, GGML_TYPE_F16)); // tiled, nd < tile depth ++ test_cases.emplace_back(new test_dsa_score(128, 4096, 24, 512, GGML_TYPE_F32)); // large-nt fallback (warp kernel) + test_cases.emplace_back(new test_fill(3.5f, GGML_TYPE_F32, { 2048, 512, 2, 2 })); + test_cases.emplace_back(new test_fill(2.0f, GGML_TYPE_F16, { 303, 207, 11, 3 })); + test_cases.emplace_back(new test_fill(-152.0f, GGML_TYPE_F16, { 800, 600, 4, 4 }));