Hugging Face's logo

HelloSun
/
openPangu-2.0-Flash-base-Q3_K_M_EPQ

Text Generation
GGUF
Mixture of Experts
expert-paging
cpu
Model card Files
xet
 Community
openPangu-2.0-Flash-base-Q3_K_M_EPQ / llama.cpp-openpangu-v0.7.0.diff
HelloSun's picture
HelloSun
Initial EPQ openPangu-2.0-Flash release: sources, patches, report and results
00d4f8c verified
Raw History Blame Contribute Delete
287 kB
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<bool>();
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 = "</think>";
+ 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</think>')
+ 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 '</think>', 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.gguf> --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/amx/mmq.cpp b/ggml/src/ggml-cpu/amx/mmq.cpp
index 9f3a744b5..2d79f714e 100644
--- a/ggml/src/ggml-cpu/amx/mmq.cpp
+++ b/ggml/src/ggml-cpu/amx/mmq.cpp
@@ -418,7 +418,7 @@ void quantize_row_q8_K_vnni(const float * RESTRICT x, void * RESTRICT vy, int64_
// Quantize these floats
const float iscale = 127.f / amax;
- y[i].d = GGML_CPU_FP32_TO_FP16(1 / iscale);
+ y[i].d = 1 / iscale;
const float id = ( amax != 0.0f ) ? iscale : 0.f;
const __m512 vscale = _mm512_set1_ps(id);
@@ -470,12 +470,7 @@ inline void from_float<block_q8_1>(const float * x, char * vy, int64_t k) {
template <>
inline void from_float<block_q8_K>(const float * x, char * vy, int64_t k) {
-#if 1
- // TODO: this is reference impl!
- quantize_row_q8_K_ref(x, (block_q8_K *)vy, k);
-#else
quantize_row_q8_K_vnni(x, vy, k);
-#endif
}
// load A from memory to array when nrows can not fill in whole tile
diff --git a/ggml/src/ggml-cpu/arch/x86/quants.c b/ggml/src/ggml-cpu/arch/x86/quants.c
index 21cd6312a..9d05aac0e 100644
--- a/ggml/src/ggml-cpu/arch/x86/quants.c
+++ b/ggml/src/ggml-cpu/arch/x86/quants.c
@@ -2302,11 +2302,107 @@ void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
const __m128i mzero = _mm_setzero_si128();
const __m256i mone = _mm256_set1_epi8(1);
- __m256 acc = _mm256_setzero_ps();
-
float summs = 0.f;
- for (int i = 0; i < nb; ++i) {
+ int i = 0;
+
+#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__)
+ __m512 acc512 = _mm512_setzero_ps();
+ const __m256i scale_shuffle = _mm256_broadcastsi128_si256(_mm_setr_epi8(
+ 0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 8, 9, 10, 11
+ ));
+ const __m256i scale_hi_shuffle = _mm256_broadcastsi128_si256(_mm_setr_epi8(
+ -1, -1, -1, -1, 0, 1, 2, 3, -1, -1, -1, -1, 4, 5, 6, 7
+ ));
+ const __m256i scale_mask = _mm256_setr_epi32(kmask1, kmask2, kmask1, kmask2, kmask1, kmask2, kmask1, kmask2);
+ const __m256i scale_hi_mask = _mm256_setr_epi32(0, kmask3, 0, kmask3, 0, kmask3, 0, kmask3);
+
+ for (; i + 1 < nb; i += 2) {
+ const __m256i scale_bytes = MM256_SET_M128I(
+ _mm_loadu_si128((const __m128i *) x[i + 1].scales),
+ _mm_loadu_si128((const __m128i *) x[i + 0].scales)
+ );
+ const __m256i scale_ordered = _mm256_shuffle_epi8(scale_bytes, scale_shuffle);
+ const __m256i scale_low = _mm256_and_si256(_mm256_blend_epi32(scale_ordered, _mm256_srli_epi16(scale_ordered, 4), 0x88), scale_mask);
+ const __m256i scale_high = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(_mm256_shuffle_epi8(scale_bytes, scale_hi_shuffle), 6), scale_hi_mask), 4);
+ const __m256i mins_and_scales8 = _mm256_or_si256(scale_low, scale_high);
+ const __m256i mins_and_scales0 = _mm256_cvtepu8_epi16(_mm256_extracti128_si256(mins_and_scales8, 0));
+ const __m256i mins_and_scales1 = _mm256_cvtepu8_epi16(_mm256_extracti128_si256(mins_and_scales8, 1));
+
+ const __m256i q8sums0 = _mm256_loadu_si256((const __m256i *) y[i + 0].bsums);
+ const __m128i q8s0 = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums0, 0), _mm256_extracti128_si256(q8sums0, 1));
+ const __m128i prod0 = _mm_madd_epi16(_mm256_extracti128_si256(mins_and_scales0, 1), q8s0);
+ const __m128i hsum0 = _mm_hadd_epi32(_mm_hadd_epi32(prod0, mzero), mzero);
+ summs -= y[i + 0].d * GGML_CPU_FP16_TO_FP32(x[i + 0].dmin) * _mm_extract_epi32(hsum0, 0);
+
+ const __m256i q8sums1 = _mm256_loadu_si256((const __m256i *) y[i + 1].bsums);
+ const __m128i q8s1 = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums1, 0), _mm256_extracti128_si256(q8sums1, 1));
+ const __m128i prod1 = _mm_madd_epi16(_mm256_extracti128_si256(mins_and_scales1, 1), q8s1);
+ const __m128i hsum1 = _mm_hadd_epi32(_mm_hadd_epi32(prod1, mzero), mzero);
+ summs -= y[i + 1].d * GGML_CPU_FP16_TO_FP32(x[i + 1].dmin) * _mm_extract_epi32(hsum1, 0);
+
+ const __m128i sc1280 = _mm256_extracti128_si256(mins_and_scales0, 0);
+ const __m128i sc1281 = _mm256_extracti128_si256(mins_and_scales1, 0);
+ const __m256i scales0 = MM256_SET_M128I(sc1280, sc1280);
+ const __m256i scales1 = MM256_SET_M128I(sc1281, sc1281);
+ const __m512i scales = _mm512_inserti64x4(_mm512_castsi256_si512(scales0), scales1, 1);
+
+ const __m256i hbits0 = _mm256_loadu_si256((const __m256i *) x[i + 0].qh);
+ const __m256i hbits1 = _mm256_loadu_si256((const __m256i *) x[i + 1].qh);
+ const __m512i hbits = _mm512_inserti64x4(_mm512_castsi256_si512(hbits0), hbits1, 1);
+ __m512i hmask = _mm512_set1_epi8(1);
+ __m512i sumi = _mm512_setzero_si512();
+
+ int bit = 0;
+
+ for (int j = 0; j < QK_K/64; ++j) {
+ const __m512i scale_0 = _mm512_shuffle_epi8(scales, _mm512_broadcast_i64x4(get_scale_shuffle_k4(2*j + 0)));
+ const __m512i scale_1 = _mm512_shuffle_epi8(scales, _mm512_broadcast_i64x4(get_scale_shuffle_k4(2*j + 1)));
+
+ const __m256i q5bits0 = _mm256_loadu_si256((const __m256i *) (x[i + 0].qs + 32*j));
+ const __m256i q5bits1 = _mm256_loadu_si256((const __m256i *) (x[i + 1].qs + 32*j));
+ const __m512i q5bits = _mm512_inserti64x4(_mm512_castsi256_si512(q5bits0), q5bits1, 1);
+
+ const __m512i q5l_0 = _mm512_and_si512(q5bits, _mm512_set1_epi8(0xF));
+ const __m512i q5h_0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask), bit++), 4);
+ const __m512i q5_0 = _mm512_add_epi8(q5l_0, q5h_0);
+ hmask = _mm512_slli_epi16(hmask, 1);
+
+ const __m512i q5l_1 = _mm512_and_si512(_mm512_srli_epi16(q5bits, 4), _mm512_set1_epi8(0xF));
+ const __m512i q5h_1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask), bit++), 4);
+ const __m512i q5_1 = _mm512_add_epi8(q5l_1, q5h_1);
+ hmask = _mm512_slli_epi16(hmask, 1);
+
+ const __m256i q8_00 = _mm256_loadu_si256((const __m256i *) (y[i + 0].qs + 64*j + 0));
+ const __m256i q8_01 = _mm256_loadu_si256((const __m256i *) (y[i + 1].qs + 64*j + 0));
+ const __m512i q8_0 = _mm512_inserti64x4(_mm512_castsi256_si512(q8_00), q8_01, 1);
+ const __m256i q8_10 = _mm256_loadu_si256((const __m256i *) (y[i + 0].qs + 64*j + 32));
+ const __m256i q8_11 = _mm256_loadu_si256((const __m256i *) (y[i + 1].qs + 64*j + 32));
+ const __m512i q8_1 = _mm512_inserti64x4(_mm512_castsi256_si512(q8_10), q8_11, 1);
+
+ const __m512i p16_0 = _mm512_maddubs_epi16(q5_0, q8_0);
+ const __m512i p16_1 = _mm512_maddubs_epi16(q5_1, q8_1);
+#if defined(__AVX512VNNI__)
+ sumi = _mm512_dpwssd_epi32(sumi, scale_0, p16_0);
+ sumi = _mm512_dpwssd_epi32(sumi, scale_1, p16_1);
+#else
+ const __m512i p32_0 = _mm512_madd_epi16(scale_0, p16_0);
+ const __m512i p32_1 = _mm512_madd_epi16(scale_1, p16_1);
+ sumi = _mm512_add_epi32(sumi, _mm512_add_epi32(p32_0, p32_1));
+#endif
+ }
+
+ __m512 vd = _mm512_castps256_ps512(_mm256_set1_ps(y[i + 0].d * GGML_CPU_FP16_TO_FP32(x[i + 0].d)));
+ vd = _mm512_insertf32x8(vd, _mm256_set1_ps(y[i + 1].d * GGML_CPU_FP16_TO_FP32(x[i + 1].d)), 1);
+ acc512 = _mm512_fmadd_ps(vd, _mm512_cvtepi32_ps(sumi), acc512);
+ }
+
+ __m256 acc = _mm256_add_ps(_mm512_castps512_ps256(acc512), _mm512_extractf32x8_ps(acc512, 1));
+#else
+ __m256 acc = _mm256_setzero_ps();
+#endif
+
+ for (; i < nb; ++i) {
const uint8_t * GGML_RESTRICT q5 = x[i].qs;
const int8_t * GGML_RESTRICT q8 = y[i].qs;
@@ -2358,13 +2454,20 @@ void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
const __m256i q8_0 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
const __m256i q8_1 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
- __m256i p16_0 = _mm256_maddubs_epi16(q5_0, q8_0);
- __m256i p16_1 = _mm256_maddubs_epi16(q5_1, q8_1);
+ const __m256i p16_0 = _mm256_maddubs_epi16(q5_0, q8_0);
+ const __m256i p16_1 = _mm256_maddubs_epi16(q5_1, q8_1);
- p16_0 = _mm256_madd_epi16(scale_0, p16_0);
- p16_1 = _mm256_madd_epi16(scale_1, p16_1);
-
- sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_0, p16_1));
+#if defined(__AVX512VNNI__) && defined(__AVX512VL__)
+ sumi = _mm256_dpwssd_epi32(sumi, scale_0, p16_0);
+ sumi = _mm256_dpwssd_epi32(sumi, scale_1, p16_1);
+#elif defined(__AVXVNNI__)
+ sumi = _mm256_dpwssd_avx_epi32(sumi, scale_0, p16_0);
+ sumi = _mm256_dpwssd_avx_epi32(sumi, scale_1, p16_1);
+#else
+ const __m256i p32_0 = _mm256_madd_epi16(scale_0, p16_0);
+ const __m256i p32_1 = _mm256_madd_epi16(scale_1, p16_1);
+ sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p32_0, p32_1));
+#endif
}
@@ -2505,9 +2608,105 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
const __m256i m3 = _mm256_set1_epi8(3);
const __m256i m15 = _mm256_set1_epi8(15);
+ int i = 0;
+
+#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__)
+ const __m512i m3_512 = _mm512_set1_epi8(3);
+ const __m512i m12_512 = _mm512_set1_epi8(12);
+ const __m512i m15_512 = _mm512_set1_epi8(15);
+ const __m512i m48_512 = _mm512_set1_epi8(48);
+ const __m512i mc0_512 = _mm512_set1_epi8(-64);
+
+ __m512 acc512 = _mm512_setzero_ps();
+
+ for (; i + 1 < nb; i += 2) {
+ const __m256i q8sums0 = _mm256_loadu_si256((const __m256i *) y[i + 0].bsums);
+ const __m256i q8sums1 = _mm256_loadu_si256((const __m256i *) y[i + 1].bsums);
+ const __m512i q8sums = _mm512_inserti64x4(_mm512_castsi256_si512(q8sums0), q8sums1, 1);
+
+ const __m128i scales0 = _mm_loadu_si128((const __m128i *) x[i + 0].scales);
+ const __m128i scales1 = _mm_loadu_si128((const __m128i *) x[i + 1].scales);
+ const __m256i scales = MM256_SET_M128I(scales1, scales0);
+ const __m512i scales_16 = _mm512_cvtepi8_epi16(scales);
+ const __m512i q8sclsub = _mm512_slli_epi32(_mm512_madd_epi16(q8sums, scales_16), 5);
+
+ __m512i sumi = _mm512_setzero_si512();
+
+ for (int j = 0; j < QK_K/128; ++j) {
+ const __m256i q4bits10 = _mm256_loadu_si256((const __m256i *) (x[i + 0].ql + 64*j + 0));
+ const __m256i q4bits11 = _mm256_loadu_si256((const __m256i *) (x[i + 1].ql + 64*j + 0));
+ const __m512i q4bits1 = _mm512_inserti64x4(_mm512_castsi256_si512(q4bits10), q4bits11, 1);
+ const __m256i q4bits20 = _mm256_loadu_si256((const __m256i *) (x[i + 0].ql + 64*j + 32));
+ const __m256i q4bits21 = _mm256_loadu_si256((const __m256i *) (x[i + 1].ql + 64*j + 32));
+ const __m512i q4bits2 = _mm512_inserti64x4(_mm512_castsi256_si512(q4bits20), q4bits21, 1);
+ const __m256i q4bitsH0 = _mm256_loadu_si256((const __m256i *) (x[i + 0].qh + 32*j));
+ const __m256i q4bitsH1 = _mm256_loadu_si256((const __m256i *) (x[i + 1].qh + 32*j));
+ const __m512i q4bitsH = _mm512_inserti64x4(_mm512_castsi256_si512(q4bitsH0), q4bitsH1, 1);
+
+ const __m512i q4h_0 = _mm512_slli_epi16(_mm512_and_si512(q4bitsH, m3_512), 4);
+ const __m512i q4h_1 = _mm512_slli_epi16(_mm512_and_si512(q4bitsH, m12_512), 2);
+ const __m512i q4h_2 = _mm512_and_si512(q4bitsH, m48_512);
+ const __m512i q4h_3 = _mm512_srli_epi16(_mm512_and_si512(q4bitsH, mc0_512), 2);
+
+ const __m512i q4_0 = _mm512_or_si512(_mm512_and_si512(q4bits1, m15_512), q4h_0);
+ const __m512i q4_1 = _mm512_or_si512(_mm512_and_si512(q4bits2, m15_512), q4h_1);
+ const __m512i q4_2 = _mm512_or_si512(_mm512_and_si512(_mm512_srli_epi16(q4bits1, 4), m15_512), q4h_2);
+ const __m512i q4_3 = _mm512_or_si512(_mm512_and_si512(_mm512_srli_epi16(q4bits2, 4), m15_512), q4h_3);
+
+ const __m256i q8_00 = _mm256_loadu_si256((const __m256i *) (y[i + 0].qs + 128*j + 0));
+ const __m256i q8_01 = _mm256_loadu_si256((const __m256i *) (y[i + 1].qs + 128*j + 0));
+ const __m512i q8_0 = _mm512_inserti64x4(_mm512_castsi256_si512(q8_00), q8_01, 1);
+ const __m256i q8_10 = _mm256_loadu_si256((const __m256i *) (y[i + 0].qs + 128*j + 32));
+ const __m256i q8_11 = _mm256_loadu_si256((const __m256i *) (y[i + 1].qs + 128*j + 32));
+ const __m512i q8_1 = _mm512_inserti64x4(_mm512_castsi256_si512(q8_10), q8_11, 1);
+ const __m256i q8_20 = _mm256_loadu_si256((const __m256i *) (y[i + 0].qs + 128*j + 64));
+ const __m256i q8_21 = _mm256_loadu_si256((const __m256i *) (y[i + 1].qs + 128*j + 64));
+ const __m512i q8_2 = _mm512_inserti64x4(_mm512_castsi256_si512(q8_20), q8_21, 1);
+ const __m256i q8_30 = _mm256_loadu_si256((const __m256i *) (y[i + 0].qs + 128*j + 96));
+ const __m256i q8_31 = _mm256_loadu_si256((const __m256i *) (y[i + 1].qs + 128*j + 96));
+ const __m512i q8_3 = _mm512_inserti64x4(_mm512_castsi256_si512(q8_30), q8_31, 1);
+
+ const __m256i scale_shuffle0 = MM256_SET_M128I(get_scale_shuffle(4*j + 0), get_scale_shuffle(4*j + 0));
+ const __m256i scale_shuffle1 = MM256_SET_M128I(get_scale_shuffle(4*j + 1), get_scale_shuffle(4*j + 1));
+ const __m256i scale_shuffle2 = MM256_SET_M128I(get_scale_shuffle(4*j + 2), get_scale_shuffle(4*j + 2));
+ const __m256i scale_shuffle3 = MM256_SET_M128I(get_scale_shuffle(4*j + 3), get_scale_shuffle(4*j + 3));
+
+ const __m512i scale_0 = _mm512_cvtepi8_epi16(_mm256_shuffle_epi8(scales, scale_shuffle0));
+ const __m512i scale_1 = _mm512_cvtepi8_epi16(_mm256_shuffle_epi8(scales, scale_shuffle1));
+ const __m512i scale_2 = _mm512_cvtepi8_epi16(_mm256_shuffle_epi8(scales, scale_shuffle2));
+ const __m512i scale_3 = _mm512_cvtepi8_epi16(_mm256_shuffle_epi8(scales, scale_shuffle3));
+ const __m512i p16_0 = _mm512_maddubs_epi16(q4_0, q8_0);
+ const __m512i p16_1 = _mm512_maddubs_epi16(q4_1, q8_1);
+ const __m512i p16_2 = _mm512_maddubs_epi16(q4_2, q8_2);
+ const __m512i p16_3 = _mm512_maddubs_epi16(q4_3, q8_3);
+
+#if defined(__AVX512VNNI__)
+ sumi = _mm512_dpwssd_epi32(sumi, scale_0, p16_0);
+ sumi = _mm512_dpwssd_epi32(sumi, scale_1, p16_1);
+ sumi = _mm512_dpwssd_epi32(sumi, scale_2, p16_2);
+ sumi = _mm512_dpwssd_epi32(sumi, scale_3, p16_3);
+#else
+ const __m512i p32_0 = _mm512_madd_epi16(scale_0, p16_0);
+ const __m512i p32_1 = _mm512_madd_epi16(scale_1, p16_1);
+ const __m512i p32_2 = _mm512_madd_epi16(scale_2, p16_2);
+ const __m512i p32_3 = _mm512_madd_epi16(scale_3, p16_3);
+ sumi = _mm512_add_epi32(sumi, _mm512_add_epi32(p32_0, p32_1));
+ sumi = _mm512_add_epi32(sumi, _mm512_add_epi32(p32_2, p32_3));
+#endif
+ }
+
+ sumi = _mm512_sub_epi32(sumi, q8sclsub);
+ __m512 vd = _mm512_castps256_ps512(_mm256_set1_ps(y[i + 0].d * GGML_CPU_FP16_TO_FP32(x[i + 0].d)));
+ vd = _mm512_insertf32x8(vd, _mm256_set1_ps(y[i + 1].d * GGML_CPU_FP16_TO_FP32(x[i + 1].d)), 1);
+ acc512 = _mm512_fmadd_ps(vd, _mm512_cvtepi32_ps(sumi), acc512);
+ }
+
+ __m256 acc = _mm256_add_ps(_mm512_castps512_ps256(acc512), _mm512_extractf32x8_ps(acc512, 1));
+#else
__m256 acc = _mm256_setzero_ps();
+#endif
- for (int i = 0; i < nb; ++i) {
+ for (; i < nb; ++i) {
const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
@@ -2544,10 +2743,10 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
const __m256i q8_2 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
const __m256i q8_3 = _mm256_loadu_si256((const __m256i*)q8); q8 += 32;
- __m256i p16_0 = _mm256_maddubs_epi16(q4_0, q8_0);
- __m256i p16_1 = _mm256_maddubs_epi16(q4_1, q8_1);
- __m256i p16_2 = _mm256_maddubs_epi16(q4_2, q8_2);
- __m256i p16_3 = _mm256_maddubs_epi16(q4_3, q8_3);
+ const __m256i p16_0 = _mm256_maddubs_epi16(q4_0, q8_0);
+ const __m256i p16_1 = _mm256_maddubs_epi16(q4_1, q8_1);
+ const __m256i p16_2 = _mm256_maddubs_epi16(q4_2, q8_2);
+ const __m256i p16_3 = _mm256_maddubs_epi16(q4_3, q8_3);
const __m128i scale_0 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 0));
const __m128i scale_1 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 1));
@@ -2555,13 +2754,29 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi
const __m128i scale_3 = _mm_shuffle_epi8(scales, get_scale_shuffle(is + 3));
is += 4;
- p16_0 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_0), p16_0);
- p16_1 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_1), p16_1);
- p16_2 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_2), p16_2);
- p16_3 = _mm256_madd_epi16(_mm256_cvtepi8_epi16(scale_3), p16_3);
+ const __m256i scale16_0 = _mm256_cvtepi8_epi16(scale_0);
+ const __m256i scale16_1 = _mm256_cvtepi8_epi16(scale_1);
+ const __m256i scale16_2 = _mm256_cvtepi8_epi16(scale_2);
+ const __m256i scale16_3 = _mm256_cvtepi8_epi16(scale_3);
- sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_0, p16_1));
- sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p16_2, p16_3));
+#if defined(__AVX512VNNI__) && defined(__AVX512VL__)
+ sumi = _mm256_dpwssd_epi32(sumi, scale16_0, p16_0);
+ sumi = _mm256_dpwssd_epi32(sumi, scale16_1, p16_1);
+ sumi = _mm256_dpwssd_epi32(sumi, scale16_2, p16_2);
+ sumi = _mm256_dpwssd_epi32(sumi, scale16_3, p16_3);
+#elif defined(__AVXVNNI__)
+ sumi = _mm256_dpwssd_avx_epi32(sumi, scale16_0, p16_0);
+ sumi = _mm256_dpwssd_avx_epi32(sumi, scale16_1, p16_1);
+ sumi = _mm256_dpwssd_avx_epi32(sumi, scale16_2, p16_2);
+ sumi = _mm256_dpwssd_avx_epi32(sumi, scale16_3, p16_3);
+#else
+ const __m256i p32_0 = _mm256_madd_epi16(scale16_0, p16_0);
+ const __m256i p32_1 = _mm256_madd_epi16(scale16_1, p16_1);
+ const __m256i p32_2 = _mm256_madd_epi16(scale16_2, p16_2);
+ const __m256i p32_3 = _mm256_madd_epi16(scale16_3, p16_3);
+ sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p32_0, p32_1));
+ sumi = _mm256_add_epi32(sumi, _mm256_add_epi32(p32_2, p32_3));
+#endif
}
diff --git a/ggml/src/ggml-cpu/arch/x86/repack.cpp b/ggml/src/ggml-cpu/arch/x86/repack.cpp
index af1cebad1..f30d822c5 100644
--- a/ggml/src/ggml-cpu/arch/x86/repack.cpp
+++ b/ggml/src/ggml-cpu/arch/x86/repack.cpp
@@ -2062,7 +2062,7 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
UNUSED(ncols_interleaved);
UNUSED(blocklen);
-#if defined(__AVX2__) || defined(__AVX512F__)
+#if defined(__AVX2__)
const block_q4_Kx8 * b_ptr_start = (const block_q4_Kx8 * ) vx;
const block_q8_Kx4 * a_ptr_start = (const block_q8_Kx4 * ) vy;
int64_t b_nb = n / QK_K;
@@ -2074,8 +2074,14 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
__m256i requiredOrder = _mm256_set_epi32(3, 2, 1, 0, 7, 6, 5, 4);
int64_t xstart = 0;
int anr = nr - nr % 16;; // Used to align nr with boundary of 16
-#if defined(__AVX512BW__) && defined(__AVX512DQ__)
+#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512DQ__)
int anc = nc - nc % 16; // Used to align nc with boundary of 16
+
+ const __m512i rhsLeftOrderSp1 = _mm512_setr_epi32(0x00, 0x02, 0x00, 0x02, 0x08, 0x0A, 0x08, 0x0A, 0x10, 0x12, 0x10, 0x12, 0x18, 0x1A, 0x18, 0x1A);
+ const __m512i rhsLeftOrderSp2 = _mm512_setr_epi32(0x01, 0x03, 0x01, 0x03, 0x09, 0x0B, 0x09, 0x0B, 0x11, 0x13, 0x11, 0x13, 0x19, 0x1B, 0x19, 0x1B);
+ const __m512i rhsRightOrderSp1 = _mm512_setr_epi32(0x04, 0x06, 0x04, 0x06, 0x0C, 0x0E, 0x0C, 0x0E, 0x14, 0x16, 0x14, 0x16, 0x1C, 0x1E, 0x1C, 0x1E);
+ const __m512i rhsRightOrderSp2 = _mm512_setr_epi32(0x05, 0x07, 0x05, 0x07, 0x0D, 0x0F, 0x0D, 0x0F, 0x15, 0x17, 0x15, 0x17, 0x1D, 0x1F, 0x1D, 0x1F);
+
// Mask to mask out nibbles from packed bytes expanded to 512 bit length
const __m512i m4bexpanded = _mm512_set1_epi8(0x0F);
//Take group of four block_q8_Kx4 structures at each pass of the loop and perform dot product operation
@@ -2115,111 +2121,85 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
// Loop to iterate over the eight sub blocks of a super block - two sub blocks are processed per iteration
for (int sb = 0; sb < QK_K / 64; sb++) {
+ const __m512i rhs_raw_mat_01234567_0 = _mm512_loadu_si512((const __m512i * )(b_ptr_0[b].qs + sb * 256));
+ const __m512i rhs_raw_mat_01234567_1 = _mm512_loadu_si512((const __m512i * )(b_ptr_0[b].qs + 64 + sb * 256));
+ const __m512i rhs_raw_mat_01234567_2 = _mm512_loadu_si512((const __m512i * )(b_ptr_0[b].qs + 128 + sb * 256));
+ const __m512i rhs_raw_mat_01234567_3 = _mm512_loadu_si512((const __m512i * )(b_ptr_0[b].qs + 192 + sb * 256));
- const __m256i rhs_raw_mat_0123_0 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + sb * 256));
- const __m256i rhs_raw_mat_4567_0 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 32 + sb * 256));
- const __m256i rhs_raw_mat_0123_1 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 64 + sb * 256));
- const __m256i rhs_raw_mat_4567_1 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 96 + sb * 256));
- const __m256i rhs_raw_mat_0123_2 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 128 + sb * 256));
- const __m256i rhs_raw_mat_4567_2 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 160 + sb * 256));
- const __m256i rhs_raw_mat_0123_3 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 192 + sb * 256));
- const __m256i rhs_raw_mat_4567_3 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 224 + sb * 256));
+ const __m512i rhs_raw_mat_89ABCDEF_0 = _mm512_loadu_si512((const __m512i * )(b_ptr_1[b].qs + sb * 256));
+ const __m512i rhs_raw_mat_89ABCDEF_1 = _mm512_loadu_si512((const __m512i * )(b_ptr_1[b].qs + 64 + sb * 256));
+ const __m512i rhs_raw_mat_89ABCDEF_2 = _mm512_loadu_si512((const __m512i * )(b_ptr_1[b].qs + 128 + sb * 256));
+ const __m512i rhs_raw_mat_89ABCDEF_3 = _mm512_loadu_si512((const __m512i * )(b_ptr_1[b].qs + 192 + sb * 256));
- const __m256i rhs_raw_mat_89AB_0 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + sb * 256));
- const __m256i rhs_raw_mat_CDEF_0 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 32 + sb * 256));
- const __m256i rhs_raw_mat_89AB_1 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 64 + sb * 256));
- const __m256i rhs_raw_mat_CDEF_1 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 96 + sb * 256));
- const __m256i rhs_raw_mat_89AB_2 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 128 + sb * 256));
- const __m256i rhs_raw_mat_CDEF_2 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 160 + sb * 256));
- const __m256i rhs_raw_mat_89AB_3 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 192 + sb * 256));
- const __m256i rhs_raw_mat_CDEF_3 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 224 + sb * 256));
+ const __m512i rhs_raw_mat_014589CD_0_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_0, rhsLeftOrderSp1, rhs_raw_mat_89ABCDEF_0);
+ const __m512i rhs_raw_mat_014589CD_0_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_0, rhsLeftOrderSp2, rhs_raw_mat_89ABCDEF_0);
+ const __m512i rhs_raw_mat_2367ABEF_0_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_0, rhsRightOrderSp1, rhs_raw_mat_89ABCDEF_0);
+ const __m512i rhs_raw_mat_2367ABEF_0_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_0, rhsRightOrderSp2, rhs_raw_mat_89ABCDEF_0);
- const __m256i rhs_raw_mat_0145_0 = _mm256_blend_epi32(rhs_raw_mat_0123_0, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_0, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_0 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_0, requiredOrder), rhs_raw_mat_4567_0, 240);
- const __m256i rhs_raw_mat_0145_1 = _mm256_blend_epi32(rhs_raw_mat_0123_1, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_1, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_1 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_1, requiredOrder), rhs_raw_mat_4567_1, 240);
- const __m256i rhs_raw_mat_0145_2 = _mm256_blend_epi32(rhs_raw_mat_0123_2, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_2, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_2 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_2, requiredOrder), rhs_raw_mat_4567_2, 240);
- const __m256i rhs_raw_mat_0145_3 = _mm256_blend_epi32(rhs_raw_mat_0123_3, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_3, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_3 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_3, requiredOrder), rhs_raw_mat_4567_3, 240);
+ const __m512i rhs_raw_mat_014589CD_1_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_1, rhsLeftOrderSp1, rhs_raw_mat_89ABCDEF_1);
+ const __m512i rhs_raw_mat_014589CD_1_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_1, rhsLeftOrderSp2, rhs_raw_mat_89ABCDEF_1);
+ const __m512i rhs_raw_mat_2367ABEF_1_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_1, rhsRightOrderSp1, rhs_raw_mat_89ABCDEF_1);
+ const __m512i rhs_raw_mat_2367ABEF_1_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_1, rhsRightOrderSp2, rhs_raw_mat_89ABCDEF_1);
- const __m256i rhs_raw_mat_89CD_0 = _mm256_blend_epi32(rhs_raw_mat_89AB_0, _mm256_permutevar8x32_epi32(rhs_raw_mat_CDEF_0, requiredOrder), 240);
- const __m256i rhs_raw_mat_ABEF_0 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_89AB_0, requiredOrder), rhs_raw_mat_CDEF_0, 240);
- const __m256i rhs_raw_mat_89CD_1 = _mm256_blend_epi32(rhs_raw_mat_89AB_1, _mm256_permutevar8x32_epi32(rhs_raw_mat_CDEF_1, requiredOrder), 240);
- const __m256i rhs_raw_mat_ABEF_1 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_89AB_1, requiredOrder), rhs_raw_mat_CDEF_1, 240);
- const __m256i rhs_raw_mat_89CD_2 = _mm256_blend_epi32(rhs_raw_mat_89AB_2, _mm256_permutevar8x32_epi32(rhs_raw_mat_CDEF_2, requiredOrder), 240);
- const __m256i rhs_raw_mat_ABEF_2 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_89AB_2, requiredOrder), rhs_raw_mat_CDEF_2, 240);
- const __m256i rhs_raw_mat_89CD_3 = _mm256_blend_epi32(rhs_raw_mat_89AB_3, _mm256_permutevar8x32_epi32(rhs_raw_mat_CDEF_3, requiredOrder), 240);
- const __m256i rhs_raw_mat_ABEF_3 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_89AB_3, requiredOrder), rhs_raw_mat_CDEF_3, 240);
+ const __m512i rhs_raw_mat_014589CD_2_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_2, rhsLeftOrderSp1, rhs_raw_mat_89ABCDEF_2);
+ const __m512i rhs_raw_mat_014589CD_2_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_2, rhsLeftOrderSp2, rhs_raw_mat_89ABCDEF_2);
+ const __m512i rhs_raw_mat_2367ABEF_2_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_2, rhsRightOrderSp1, rhs_raw_mat_89ABCDEF_2);
+ const __m512i rhs_raw_mat_2367ABEF_2_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_2, rhsRightOrderSp2, rhs_raw_mat_89ABCDEF_2);
- const __m512i rhs_raw_mat_014589CD_0 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_0145_0), rhs_raw_mat_89CD_0, 1);
- const __m512i rhs_raw_mat_2367ABEF_0 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_2367_0), rhs_raw_mat_ABEF_0, 1);
- const __m512i rhs_raw_mat_014589CD_1 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_0145_1), rhs_raw_mat_89CD_1, 1);
- const __m512i rhs_raw_mat_2367ABEF_1 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_2367_1), rhs_raw_mat_ABEF_1, 1);
+ const __m512i rhs_raw_mat_014589CD_3_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_3, rhsLeftOrderSp1, rhs_raw_mat_89ABCDEF_3);
+ const __m512i rhs_raw_mat_014589CD_3_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_3, rhsLeftOrderSp2, rhs_raw_mat_89ABCDEF_3);
+ const __m512i rhs_raw_mat_2367ABEF_3_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_3, rhsRightOrderSp1, rhs_raw_mat_89ABCDEF_3);
+ const __m512i rhs_raw_mat_2367ABEF_3_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_3, rhsRightOrderSp2, rhs_raw_mat_89ABCDEF_3);
- const __m512i rhs_raw_mat_014589CD_2 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_0145_2), rhs_raw_mat_89CD_2, 1);
- const __m512i rhs_raw_mat_2367ABEF_2 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_2367_2), rhs_raw_mat_ABEF_2, 1);
- const __m512i rhs_raw_mat_014589CD_3 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_0145_3), rhs_raw_mat_89CD_3, 1);
- const __m512i rhs_raw_mat_2367ABEF_3 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_2367_3), rhs_raw_mat_ABEF_3, 1);
+ // 4bit-8bit pattern one - right side input
+ const __m512i rhs_mat_014589CD_00_sp1 = _mm512_and_si512(rhs_raw_mat_014589CD_0_sp1, m4bexpanded); //B00(0-3) B01(0-3) B00(0-3) B01(0-3) B04(0-3) B05(0-3) B04(0-3) B05(0-3) B08(0-3) B09(0-3) B08(0-3) B09(0-3) B0C(0-3) B0D(0-3) B0C(0-3) B0D(0-3)
+ const __m512i rhs_mat_014589CD_10_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_0_sp1, 4), m4bexpanded); //B10(0-3) B11(0-3) B10(0-3) B11(0-3) B14(0-3) B15(0-3) B14(0-3) B15(0-3) B18(0-3) B19(0-3) B18(0-3) B19(0-3) B1C(0-3) B1D(0-3) B1C(0-3) B1D(0-3)
- //4-bit -> 8-bit
- const __m512i rhs_mat_014589CD_00 = _mm512_and_si512(rhs_raw_mat_014589CD_0, m4bexpanded); //B00(0-7) B01(0-7) B04(0-7) B05(0-7) B08(0-7) B09(0-7) B0C(0-7) B0D(0-7)
- const __m512i rhs_mat_2367ABEF_00 = _mm512_and_si512(rhs_raw_mat_2367ABEF_0, m4bexpanded); //B02(0-7) B03(0-7) B06(0-7) B07(0-7) B0A(0-7) B0B(0-7) B0E(0-7) B0F(0-7)
- const __m512i rhs_mat_014589CD_01 = _mm512_and_si512(rhs_raw_mat_014589CD_1, m4bexpanded); //B00(8-15) B01(8-15) B04(8-15) B05(8-15) B08(8-15) B09(8-15) B0C(8-15) B0D(8-15)
- const __m512i rhs_mat_2367ABEF_01 = _mm512_and_si512(rhs_raw_mat_2367ABEF_1, m4bexpanded); //B02(8-15) B03(8-15) B06(8-15) B07(8-15) B0A(8-15) B0B(8-15) B0E(8-15) B0F(8-15)
+ const __m512i rhs_mat_2367ABEF_00_sp1 = _mm512_and_si512(rhs_raw_mat_2367ABEF_0_sp1, m4bexpanded); //B02(0-3) B03(0-3) B02(0-3) B03(0-3) B06(0-3) B07(0-3) B06(0-3) B07(0-3) B0A(0-3) B0B(0-3) B0A(0-3) B0B(0-3) B0E(0-3) B0F(0-3) B0E(0-3) B0F(0-3)
+ const __m512i rhs_mat_2367ABEF_10_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_0_sp1, 4), m4bexpanded); //B12(0-3) B13(0-3) B12(0-3) B13(0-3) B16(0-3) B17(0-3) B16(0-3) B17(0-3) B1A(0-3) B1B(0-3) B1A(0-3) B1B(0-3) B1E(0-3) B1F(0-3) B1E(0-3) B1F(0-3)
- const __m512i rhs_mat_014589CD_02 = _mm512_and_si512(rhs_raw_mat_014589CD_2, m4bexpanded); //B00(16-23) B01(16-23) B04(16-23) B05(16-23) B08(16-23) B09(16-23) B0C(16-23) B0D(16-23)
- const __m512i rhs_mat_2367ABEF_02 = _mm512_and_si512(rhs_raw_mat_2367ABEF_2, m4bexpanded); //B02(16-23) B03(16-23) B06(16-23) B07(16-23) B0A(16-23) B0B(16-23) B0E(16-23) B0F(16-23)
- const __m512i rhs_mat_014589CD_03 = _mm512_and_si512(rhs_raw_mat_014589CD_3, m4bexpanded); //B00(24-31) B01(24-31) B04(24-31) B05(24-31) B08(24-31) B09(24-31) B0C(24-31) B0D(24-31)
- const __m512i rhs_mat_2367ABEF_03 = _mm512_and_si512(rhs_raw_mat_2367ABEF_3, m4bexpanded); //B02(24-31) B03(24-31) B06(24-31) B07(24-31) B0A(24-31) B0B(24-31) B0E(24-31) B0F(24-31)
+ const __m512i rhs_mat_014589CD_01_sp1 = _mm512_and_si512(rhs_raw_mat_014589CD_1_sp1, m4bexpanded); //B00(8-11) B01(8-11) B00(8-11) B01(8-11) B04(8-11) B05(8-11) B04(8-11) B05(8-11) B08(8-11) B09(8-11) B08(8-11) B09(8-11) B0C(8-11) B0D(8-11) B0C(8-11) B0D(8-11)
+ const __m512i rhs_mat_014589CD_11_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_1_sp1, 4), m4bexpanded); //B10(8-11) B11(8-11) B10(8-11) B11(8-11) B14(8-11) B15(8-11) B14(8-11) B15(8-11) B18(8-11) B19(8-11) B18(8-11) B19(8-11) B1C(8-11) B1D(8-11) B1C(8-11) B1D(8-11)
- const __m512i rhs_mat_014589CD_10 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_0, 4), m4bexpanded); //B10(0-7) B11(0-7) B14(0-7) B15(0-7) B18(0-7) B19(0-7) B1C(0-7) B1D(0-7)
- const __m512i rhs_mat_2367ABEF_10 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_0, 4), m4bexpanded); //B12(0-7) B13(0-7) B16(0-7) B17(0-7) B1A(0-7) B1B(0-7) B1E(0-7) B1F(0-7)
- const __m512i rhs_mat_014589CD_11 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_1, 4), m4bexpanded); //B10(8-15) B11(8-15) B14(8-15) B15(8-15) B18(8-15) B19(8-15) B1C(8-15) B1D(8-15)
- const __m512i rhs_mat_2367ABEF_11 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_1, 4), m4bexpanded); //B12(8-15) B13(8-15) B16(8-15) B17(8-15) B1A(8-15) B1B(8-15) B1E(8-15) B1F(8-15)
+ const __m512i rhs_mat_2367ABEF_01_sp1 = _mm512_and_si512(rhs_raw_mat_2367ABEF_1_sp1, m4bexpanded); //B02(8-11) B03(8-11) B02(8-11) B03(8-11) B06(8-11) B07(8-11) B06(8-11) B07(8-11) B0A(8-11) B0B(8-11) B0A(8-11) B0B(8-11) B0E(8-11) B0F(8-11) B0E(8-11) B0F(8-11)
+ const __m512i rhs_mat_2367ABEF_11_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_1_sp1, 4), m4bexpanded); //B12(8-11) B13(8-11) B12(8-11) B13(8-11) B16(8-11) B17(8-11) B16(8-11) B17(8-11) B1A(8-11) B1B(8-11) B1A(8-11) B1B(8-1１) B１E(8-１１) B１F(8-１１) B１E(8-１１) B１F(8-１１)
- const __m512i rhs_mat_014589CD_12 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_2, 4), m4bexpanded); //B10(16-23) B11(16-23) B14(16-23) B15(16-23) B18(16-23) B19(16-23) B1C(16-23) B1D(16-23)
- const __m512i rhs_mat_2367ABEF_12 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_2, 4), m4bexpanded); //B12(16-23) B13(16-23) B16(16-23) B17(16-23) B1A(16-23) B1B(16-23) B1E(16-23) B1F(16-23)
- const __m512i rhs_mat_014589CD_13 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_3, 4), m4bexpanded); //B10(24-31) B11(24-31) B14(24-31) B15(24-31) B18(24-31) B19(24-31) B1C(24-31) B1D(24-31)
- const __m512i rhs_mat_2367ABEF_13 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_3, 4), m4bexpanded); //B12(24-31) B13(24-31) B16(24-31) B17(24-31) B1A(24-31) B1B(24-31) B1E(24-31) B1F(24-31)
+ const __m512i rhs_mat_014589CD_02_sp1 = _mm512_and_si512(rhs_raw_mat_014589CD_2_sp1, m4bexpanded); //B00(16-19) B01(16-19) B00(16-19) B01(16-19) B04(16-19) B05(16-19) B04(16-19) B05(16-19) B08(16-19) B09(16-19) B08(16-19) B09(16-19) B0C(16-19) B0D(16-19) B0C(16-19) B0D(16-19)
+ const __m512i rhs_mat_014589CD_12_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_2_sp1, 4), m4bexpanded); //B10(16-19) B11(16-19) B10(16-19) B11(16-19) B14(16-19) B15(16-19) B14(16-19) B15(16-19) B18(16-19) B19(16-19) B18(16-19) B19(16-19) B1C(16-19) B1D(16-19) B1C(16-19) B1D(16-19)
- // Shuffle pattern one - right side input
- const __m512i rhs_mat_014589CD_00_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_00, (_MM_PERM_ENUM)136); //B00(0-3) B01(0-3) B00(0-3) B01(0-3) B04(0-3) B05(0-3) B04(0-3) B05(0-3) B08(0-3) B09(0-3) B08(0-3) B09(0-3) B0C(0-3) B0D(0-3) B0C(0-3) B0D(0-3)
- const __m512i rhs_mat_2367ABEF_00_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_00, (_MM_PERM_ENUM)136); //B02(0-3) B03(0-3) B02(0-3) B03(0-3) B06(0-3) B07(0-3) B06(0-3) B07(0-3) B0A(0-3) B0B(0-3) B0A(0-3) B0B(0-3) B0E(0-3) B0F(0-3) B0E(0-3) B0F(0-3)
- const __m512i rhs_mat_014589CD_01_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_01, (_MM_PERM_ENUM)136); //B00(8-11) B01(8-11) B00(8-11) B01(8-11) B04(8-11) B05(8-11) B04(8-11) B05(8-11) B08(8-11) B09(8-11) B08(8-11) B09(8-11) B0C(8-11) B0D(8-11) B0C(8-11) B0D(8-11)
- const __m512i rhs_mat_2367ABEF_01_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_01, (_MM_PERM_ENUM)136); //B02(8-11) B03(8-11) B02(8-11) B03(8-11) B06(8-11) B07(8-11) B06(8-11) B07(8-11) B0A(8-11) B0B(8-11) B0A(8-11) B0B(8-11) B0E(8-11) B0F(8-11) B0E(8-11) B0F(8-11)
- const __m512i rhs_mat_014589CD_02_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_02, (_MM_PERM_ENUM)136); //B00(16-19) B01(16-19) B00(16-19) B01(16-19) B04(16-19) B05(16-19) B04(16-19) B05(16-19) B08(16-19) B09(16-19) B08(16-19) B09(16-19) B0C(16-19) B0D(16-19) B0C(16-19) B0D(16-19)
- const __m512i rhs_mat_2367ABEF_02_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_02, (_MM_PERM_ENUM)136); //B02(16-19) B03(16-19) B02(16-19) B03(16-19) B06(16-19) B07(16-19) B06(16-19) B07(16-19) B0A(16-19) B0B(16-19) B0A(16-19) B0B(16-19) B0E(16-19) B0F(16-19) B0E(16-19) B0F(16-19)
- const __m512i rhs_mat_014589CD_03_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_03, (_MM_PERM_ENUM)136); //B00(24-27) B01(24-27) B00(24-27) B01(24-27) B04(24-27) B05(24-27) B04(24-27) B05(24-27) B08(24-27) B09(24-27) B08(24-27) B09(24-27) B0C(24-27) B0D(24-27) B0C(24-27) B0D(24-27)
- const __m512i rhs_mat_2367ABEF_03_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_03, (_MM_PERM_ENUM)136); //B02(24-27) B03(24-27) B02(24-27) B03(24-27) B06(24-27) B07(24-27) B06(24-27) B07(24-27) B0A(24-27) B0B(24-27) B0A(24-27) B0B(24-27) B0E(24-27) B0F(24-27) B0E(24-27) B0F(24-27)
+ const __m512i rhs_mat_2367ABEF_02_sp1 = _mm512_and_si512(rhs_raw_mat_2367ABEF_2_sp1, m4bexpanded); //B02(16-19) B03(16-19) B02(16-19) B03(16-19) B06(16-19) B07(16-19) B06(16-19) B07(16-19) B0A(16-19) B0B(16-19) B0A(16-19) B0B(16-19) B0E(16-19) B0F(16-19) B0E(16-19) B0F(16-19)
+ const __m512i rhs_mat_2367ABEF_12_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_2_sp1, 4), m4bexpanded); //B12(16-19) B13(16-19) B12(16-19) B13(16-19) B16(16-19) B17(16-19) B16(16-19) B17(16-19) B1A(16-19) B1B(16-19) B1A(16-19) B1B(16-19) B1E(16-19) B1F(16-19) B1E(16-19) B1F(16-19)
- const __m512i rhs_mat_014589CD_10_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_10, (_MM_PERM_ENUM)136); //B10(0-3) B11(0-3) B10(0-3) B11(0-3) B14(0-3) B15(0-3) B14(0-3) B15(0-3) B18(0-3) B19(0-3) B18(0-3) B19(0-3) B1C(0-3) B1D(0-3) B1C(0-3) B1D(0-3)
- const __m512i rhs_mat_2367ABEF_10_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_10, (_MM_PERM_ENUM)136); //B12(0-3) B13(0-3) B12(0-3) B13(0-3) B16(0-3) B17(0-3) B16(0-3) B17(0-3) B1A(0-3) B1B(0-3) B1A(0-3) B1B(0-3) B1E(0-3) B1F(0-3) B1E(0-3) B1F(0-3)
- const __m512i rhs_mat_014589CD_11_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_11, (_MM_PERM_ENUM)136); //B10(8-11) B11(8-11) B10(8-11) B11(8-11) B14(8-11) B15(8-11) B14(8-11) B15(8-11) B18(8-11) B19(8-11) B18(8-11) B19(8-11) B1C(8-11) B1D(8-11) B1C(8-11) B1D(8-11)
- const __m512i rhs_mat_2367ABEF_11_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_11, (_MM_PERM_ENUM)136); //B12(8-11) B13(8-11) B12(8-11) B13(8-11) B16(8-11) B17(8-11) B16(8-11) B17(8-11) B1A(8-11) B1B(8-11) B1A(8-11) B1B(8-11) B1E(8-11) B1F(8-11) B1E(8-11) B1F(8-11)
- const __m512i rhs_mat_014589CD_12_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_12, (_MM_PERM_ENUM)136); //B10(16-19) B11(16-19) B10(16-19) B11(16-19) B14(16-19) B15(16-19) B14(16-19) B15(16-19) B18(16-19) B19(16-19) B18(16-19) B19(16-19) B1C(16-19) B1D(16-19) B1C(16-19) B1D(16-19)
- const __m512i rhs_mat_2367ABEF_12_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_12, (_MM_PERM_ENUM)136); //B12(16-19) B13(16-19) B12(16-19) B13(16-19) B16(16-19) B17(16-19) B16(16-19) B17(16-19) B1A(16-19) B1B(16-19) B1A(16-19) B1B(16-19) B1E(16-19) B1F(16-19) B1E(16-19) B1F(16-19)
- const __m512i rhs_mat_014589CD_13_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_13, (_MM_PERM_ENUM)136); //B10(24-27) B11(24-27) B10(24-27) B11(24-27) B14(24-27) B15(24-27) B14(24-27) B15(24-27) B18(24-27) B19(24-27) B18(24-27) B19(24-27) B1C(24-27) B1D(24-27) B1C(24-27) B1D(24-27)
- const __m512i rhs_mat_2367ABEF_13_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_13, (_MM_PERM_ENUM)136); //B12(24-27) B13(24-27) B12(24-27) B13(24-27) B16(24-27) B17(24-27) B16(24-27) B17(24-27) B1A(24-27) B1B(24-27) B1A(24-27) B1B(24-27) B1E(24-27) B1F(24-27) B1E(24-27) B1F(24-27)
+ const __m512i rhs_mat_014589CD_03_sp1 = _mm512_and_si512(rhs_raw_mat_014589CD_3_sp1, m4bexpanded); //B00(24-27) B01(24-27) B00(24-27) B01(24-27) B04(24-27) B05(24-27) B04(24-27) B05(24-27) B08(24-27) B09(24-27) B08(24-27) B09(24-27) B0C(24-27) B0D(24-27) B0C(24-27) B0D(24-27)
+ const __m512i rhs_mat_014589CD_13_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_3_sp1, 4), m4bexpanded); //B10(24-27) B11(24-27) B10(24-27) B11(24-27) B14(24-27) B15(24-27) B14(24-27) B15(24-27) B18(24-27) B19(24-27) B18(24-27) B19(24-27) B1C(24-27) B1D(24-27) B1C(24-27) B1D(24-27)
- // Shuffle pattern two - right side input
- const __m512i rhs_mat_014589CD_00_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_00, (_MM_PERM_ENUM)221); //B00(4-7) B01(4-7) B00(4-7) B01(4-7) B04(4-7) B05(4-7) B04(4-7) B05(4-7) B08(4-7) B09(4-7) B08(4-7) B09(4-7) B0C(4-7) B0D(4-7) B0C(4-7) B0D(4-7)
- const __m512i rhs_mat_2367ABEF_00_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_00, (_MM_PERM_ENUM)221); //B02(4-7) B03(4-7) B02(4-7) B03(4-7) B06(4-7) B07(4-7) B06(4-7) B07(4-7) B0A(4-7) B0B(4-7) B0A(4-7) B0B(4-7) B0E(4-7) B0F(4-7) B0E(4-7) B0F(4-7)
- const __m512i rhs_mat_014589CD_01_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_01, (_MM_PERM_ENUM)221); //B00(12-15) B01(12-15) B00(12-15) B01(12-15) B04(12-15) B05(12-15) B04(12-15) B05(12-15) B08(12-15) B09(12-15) B08(12-15) B09(12-15) B0C(12-15) B0D(12-15) B0C(12-15) B0D(12-15)
- const __m512i rhs_mat_2367ABEF_01_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_01, (_MM_PERM_ENUM)221); //B02(12-15) B03(12-15) B02(12-15) B03(12-15) B06(12-15) B07(12-15) B06(12-15) B07(12-15) B0A(12-15) B0B(12-15) B0A(12-15) B0B(12-15) B0E(12-15) B0F(12-15) B0E(12-15) B0F(12-15)
- const __m512i rhs_mat_014589CD_02_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_02, (_MM_PERM_ENUM)221); //B00(20-23) B01(20-23) B00(20-23) B01(20-23) B04(20-23) B05(20-23) B04(20-23) B05(20-23) B08(20-23) B09(20-23) B08(20-23) B09(20-23) B0C(20-23) B0D(20-23) B0C(20-23) B0D(20-23)
- const __m512i rhs_mat_2367ABEF_02_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_02, (_MM_PERM_ENUM)221); //B02(20-23) B03(20-23) B02(20-23) B03(20-23) B06(20-23) B07(20-23) B06(20-23) B07(20-23) B0A(20-23) B0B(20-23) B0A(20-23) B0B(20-23) B0E(20-23) B0F(20-23) B0E(20-23) B0F(20-23)
- const __m512i rhs_mat_014589CD_03_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_03, (_MM_PERM_ENUM)221); //B00(28-31) B01(28-31) B00(28-31) B01(28-31) B04(28-31) B05(28-31) B04(28-31) B05(28-31) B08(28-31) B09(28-31) B08(28-31) B09(28-31) B0C(28-31) B0D(28-31) B0C(28-31) 0BD(28-31)
- const __m512i rhs_mat_2367ABEF_03_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_03, (_MM_PERM_ENUM)221); //B02(28-31) B03(28-31) B02(28-31) B03(28-31) B06(28-31) B07(28-31) B06(28-31) B07(28-31) B0A(28-31) B0B(28-31) B0A(28-31) B0B(28-31) B0E(28-31) B0F(28-31) B0E(28-31) B0F(28-31)
+ const __m512i rhs_mat_2367ABEF_03_sp1 = _mm512_and_si512(rhs_raw_mat_2367ABEF_3_sp1, m4bexpanded); //B02(24-27) B03(24-27) B02(24-27) B03(24-27) B06(24-27) B07(24-27) B06(24-27) B07(24-27) B0A(24-27) B0B(24-27) B0A(24-27) B0B(24-27) B0E(24-27) B0F(24-27) B0E(24-27) B0F(24-27)
+ const __m512i rhs_mat_2367ABEF_13_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_3_sp1, 4), m4bexpanded); //B12(24-27) B13(24-27) B12(24-27) B13(24-27) B16(24-27) B17(24-27) B16(24-27) B17(24-27) B1A(24-27) B1B(24-27) B1A(24-27) B1B(24-27) B1E(24-27) B1F(24-27) B1E(24-27) B1F(24-27)
- const __m512i rhs_mat_014589CD_10_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_10, (_MM_PERM_ENUM)221); //B10(4-7) B11(4-7) B10(4-7) B11(4-7) B14(4-7) B15(4-7) B14(4-7) B15(4-7) B18(4-7) B19(4-7) B18(4-7) B19(4-7) B1C(4-7) B1D(4-7) B1C(4-7) B1D(4-7)
- const __m512i rhs_mat_2367ABEF_10_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_10, (_MM_PERM_ENUM)221); //B12(4-7) B13(4-7) B12(4-7) B13(4-7) B16(4-7) B17(4-7) B16(4-7) B17(4-7) B1A(4-7) B1B(4-7) B1A(4-7) B1B(4-7) B1E(4-7) B1F(4-7) B1E(4-7) B1F(4-7)
- const __m512i rhs_mat_014589CD_11_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_11, (_MM_PERM_ENUM)221); //B10(12-15) B11(12-15) B10(12-15) B11(12-15) B14(12-15) B15(12-15) B14(12-15) B15(12-15) B18(12-15) B19(12-15) B18(12-15) B19(12-15) B1C(12-15) B1D(12-15) B1C(12-15) B1D(12-15)
- const __m512i rhs_mat_2367ABEF_11_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_11, (_MM_PERM_ENUM)221); //B12(12-15) B13(12-15) B12(12-15) B13(12-15) B16(12-15) B17(12-15) B16(12-15) B17(12-15) B1A(12-15) B1B(12-15) B1A(12-15) B1B(12-15) B1E(12-15) B1F(12-15) B1E(12-15) B1F(12-15)
- const __m512i rhs_mat_014589CD_12_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_12, (_MM_PERM_ENUM)221); //B10(20-23) B11(20-23) B10(20-23) B11(20-23) B14(20-23) B15(20-23) B14(20-23) B15(20-23) B18(20-23) B19(20-23) B18(20-23) B19(20-23) B1C(20-23) B1D(20-23) B1C(20-23) B1D(20-23)
- const __m512i rhs_mat_2367ABEF_12_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_12, (_MM_PERM_ENUM)221); //B12(20-23) B13(20-23) B12(20-23) B13(20-23) B16(20-23) B17(20-23) B16(20-23) B17(20-23) B1A(20-23) B1B(20-23) B1A(20-23) B1B(20-23) B1E(20-23) B1F(20-23) B1E(20-23) B1F(20-23)
- const __m512i rhs_mat_014589CD_13_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_13, (_MM_PERM_ENUM)221); //B10(28-31) B11(28-31) B10(28-31) B11(28-31) B14(28-31) B15(28-31) B14(28-31) B15(28-31) B18(28-31) B19(28-31) B18(28-31) B19(28-31) B1C(28-31) B1D(28-31) B1C(28-31) B1D(28-31)
- const __m512i rhs_mat_2367ABEF_13_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_13, (_MM_PERM_ENUM)221); //B12(28-31) B13(28-31) B12(28-31) B13(28-31) B16(28-31) B17(28-31) B16(28-31) B17(28-31) B1A(28-31) B1B(28-31) B1A(28-31) B1B(28-31) B1E(28-31) B1F(28-31) B1E(28-31) B1F(28-31)
+ // 4bit-8bit pattern two - right side input
+ const __m512i rhs_mat_014589CD_00_sp2 = _mm512_and_si512(rhs_raw_mat_014589CD_0_sp2, m4bexpanded); //B00(4-7) B01(4-7) B00(4-7) B01(4-7) B04(4-7) B05(4-7) B04(4-7) B05(4-7) B08(4-7) B09(4-7) B08(4-7) B09(4-7) B0C(4-7) B0D(4-7) B0C(4-7) B0D(4-7)
+ const __m512i rhs_mat_014589CD_10_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_0_sp2, 4), m4bexpanded); //B10(4-7) B11(4-7) B10(4-7) B11(4-7) B14(4-7) B15(4-7) B14(4-7) B15(4-7) B18(4-7) B19(4-7) B18(4-7) B19(4-7) B1C(4-7) B1D(4-7) B1C(4-7) B1D(4-7)
+
+ const __m512i rhs_mat_2367ABEF_00_sp2 = _mm512_and_si512(rhs_raw_mat_2367ABEF_0_sp2, m4bexpanded); //B02(4-7) B03(4-7) B02(4-7) B03(4-7) B06(4-7) B07(4-7) B06(4-7) B07(4-7) B0A(4-7) B0B(4-7) B0A(4-7) B0B(4-7) B0E(4-7) B0F(4-7) B0E(4-7] B0F(4-7)
+ const __m512i rhs_mat_2367ABEF_10_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_0_sp2, 4), m4bexpanded); //B12(4-7] B13(4-7] B12(4-7] B13(4-7] B16(4-7] B17(4-七年] B16(四年] B1七七年] B１A（四年]B１B（四年]B１A（四年]B１B（四年]B１E（四年]B１F（四年]B１E（四年]B１F（四年]
+
+ const __m512i rhs_mat_014589CD_01_sp2 = _mm512_and_si512(rhs_raw_mat_014589CD_1_sp2, m4bexpanded); //B00(12-15) B01(12-15) B00(12-15) B01(12-15) B04(12-15) B05(12-15) B04(12-15) B05(12-15) B08(12-15) B09(12-15) B08(12-15) B09(12-15) B0C(12-15) B0D(12-15) B0C(12-15) B0D(12-15)
+ const __m512i rhs_mat_014589CD_11_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_1_sp2, 4), m4bexpanded); //B10(12-15) B11(12-15) B10(12-15) B11(12-15) B14(12-15) B15(12-15) B14(12-15) B15(12-15) B18(12-15) B19(12-15) B18(12-15) B19(12-15) B1C(12-15) B1D(12-15) B1C(12-15) B1D(12-15)
+
+ const __m512i rhs_mat_2367ABEF_01_sp2 = _mm512_and_si512(rhs_raw_mat_2367ABEF_1_sp2, m4bexpanded); //B02(12-15) B03(12-15) B02(12-15) B03(12-15) B06(12-15) B07(12-15) B06(12-15) B07(12-15) B0A(12-15) B0B(12-15) B0A(12-15) B0B(12-15) B0E(12-15) B0F(12-15) B0E(12-15) B0F(12-15)
+ const __m512i rhs_mat_2367ABEF_11_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_1_sp2, 4), m4bexpanded); //B12(12-15) B13(12-15) B12(12-15) B13(12-15) B16(12-15) B17(12-15) B16(12-15) B17(12-15) B1A(12-15) B1B(12-15) B1A(12-15) B1B(12-15) B1E(12-15) B1F(12-15) B1E(12-15) B1F(12-15)
+
+ const __m512i rhs_mat_014589CD_02_sp2 = _mm512_and_si512(rhs_raw_mat_014589CD_2_sp2, m4bexpanded); //B00(20-23) B01(20-23) B00(20-23) B01(20-23) B04(20-23) B05(20-23) B04(20-23) B05(20-23) B08(20-23) B09(20-23) B08(20-23) B09(20-23) B0C(20-23) B0D(20-23) B0C(20-23) B0D(20-23)
+ const __m512i rhs_mat_014589CD_12_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_2_sp2, 4), m4bexpanded); //B10(20-23) B11(20-23) B10(20-23) B11(20-23) B14(20-23) B15(20-23) B14(20-23) B15(20-23) B18(20-23) B19(20-23) B18(20-23) B19(20-23) B1C(20-23) B1D(20-23) B1C(20-23) B1D(20-23)
+
+ const __m512i rhs_mat_2367ABEF_02_sp2 = _mm512_and_si512(rhs_raw_mat_2367ABEF_2_sp2, m4bexpanded); //B02(20-23) B03(20-23) B02(20-23) B03(20-23) B06(20-23) B07(20-23) B06(20-23) B07(20-23) B0A(20-23) B0B(20-23) B0A(20-23) B0B(20-23) B0E(20-23) B0F(20-23) B0E(20-23) B0F(20-23)
+ const __m512i rhs_mat_2367ABEF_12_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_2_sp2, 4), m4bexpanded); //B12(20-23) B13(20-23) B12(20-23) B13(20-23) B16(20-23) B17(20-23) B16(20-23) B17(20-23) B1A(20-23) B1B(20-23) B1A(20-23) B1B(20-23) B1E(20-23) B1F(20-23) B1E(20-23) B1F(20-23)
+
+ const __m512i rhs_mat_014589CD_03_sp2 = _mm512_and_si512(rhs_raw_mat_014589CD_3_sp2, m4bexpanded); //B00(28-31) B01(28-31) B00(28-31) B01(28-31) B04(28-31) B05(28-31) B04(28-31) B05(28-31) B08(28-31) B09(28-31) B08(28-31) B09(28-31) B0C(28-31) B0D(28-31) B0C(28-31) 0BD(28-31)
+ const __m512i rhs_mat_014589CD_13_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_3_sp2, 4), m4bexpanded); //B10(28-31) B11(28-31) B10(28-31) B11(28-31) B14(28-31) B15(28-31) B14(28-31) B15(28-31) B18(28-31) B19(28-31) B18(28-31) B19(28-31) B1C(28-31) B1D(28-31) B1C(28-31) B1D(28-31)
+
+ const __m512i rhs_mat_2367ABEF_03_sp2 = _mm512_and_si512(rhs_raw_mat_2367ABEF_3_sp2, m4bexpanded); //B02(28-31) B03(28-31) B02(28-31) B03(28-31) B06(28-31) B07(28-31) B06(28-31) B07(28-31) B0A(28-31) B0B(28-31) B0A(28-31) B0B(28-31) B0E(28-31) B0F(28-31) B0E(28-31) B0F(28-31)
+ const __m512i rhs_mat_2367ABEF_13_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_3_sp2, 4), m4bexpanded); //B12(28-31) B13(28-31) B12(28-31) B13(28-31) B16(28-31) B17(28-31) B16(28-31) B17(28-31) B1A(28-31) B1B(28-31) B1A(28-31) B1B(28-31) B1E(28-31) B1F(28-31) B1E(28-31) B1F(28-31)
uint32_t utmp_00[4], utmp_01[4], utmp_10[4], utmp_11[4];
@@ -2273,51 +2253,32 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
const __m512i scale_2367ABEF_1 = _mm512_shuffle_epi32(scales_1, (_MM_PERM_ENUM)238);
for (int rp = 0; rp < 4; rp++) {
-
// Load the four block_q8_k quantized values interleaved with each other in chunks of eight bytes - A0,A1,A2,A3
- // Loaded as set of 128 bit vectors and repeated and stored into a 256 bit vector before again repeating into 512 bit vector
- __m256i lhs_mat_ymm_0123_00 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].qs + 256 * sb)));
- __m256i lhs_mat_ymm_01_00 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_00, lhs_mat_ymm_0123_00, 0);
- __m256i lhs_mat_ymm_23_00 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_00, lhs_mat_ymm_0123_00, 17);
- __m256i lhs_mat_ymm_0123_01 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].qs + 32 + 256 * sb)));
- __m256i lhs_mat_ymm_01_01 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_01, lhs_mat_ymm_0123_01, 0);
- __m256i lhs_mat_ymm_23_01 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_01, lhs_mat_ymm_0123_01, 17);
- __m256i lhs_mat_ymm_0123_02 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].qs + 64 + 256 * sb)));
- __m256i lhs_mat_ymm_01_02 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_02, lhs_mat_ymm_0123_02, 0);
- __m256i lhs_mat_ymm_23_02 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_02, lhs_mat_ymm_0123_02, 17);
- __m256i lhs_mat_ymm_0123_03 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].qs + 96 + 256 * sb)));
- __m256i lhs_mat_ymm_01_03 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_03, lhs_mat_ymm_0123_03, 0);
- __m256i lhs_mat_ymm_23_03 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_03, lhs_mat_ymm_0123_03, 17);
- __m256i lhs_mat_ymm_0123_10 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].qs + 128 + 256 * sb)));
- __m256i lhs_mat_ymm_01_10 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_10, lhs_mat_ymm_0123_10, 0);
- __m256i lhs_mat_ymm_23_10 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_10, lhs_mat_ymm_0123_10, 17);
- __m256i lhs_mat_ymm_0123_11 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].qs + 160 + 256 * sb)));
- __m256i lhs_mat_ymm_01_11 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_11, lhs_mat_ymm_0123_11, 0);
- __m256i lhs_mat_ymm_23_11 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_11, lhs_mat_ymm_0123_11, 17);
- __m256i lhs_mat_ymm_0123_12 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].qs + 192 + 256 * sb)));
- __m256i lhs_mat_ymm_01_12 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_12, lhs_mat_ymm_0123_12, 0);
- __m256i lhs_mat_ymm_23_12 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_12, lhs_mat_ymm_0123_12, 17);
- __m256i lhs_mat_ymm_0123_13 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].qs + 224 + 256 * sb)));
- __m256i lhs_mat_ymm_01_13 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_13, lhs_mat_ymm_0123_13, 0);
- __m256i lhs_mat_ymm_23_13 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_13, lhs_mat_ymm_0123_13, 17);
-
- __m512i lhs_mat_01_00 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_00), lhs_mat_ymm_01_00, 1);
- __m512i lhs_mat_23_00 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_00), lhs_mat_ymm_23_00, 1);
- __m512i lhs_mat_01_01 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_01), lhs_mat_ymm_01_01, 1);
- __m512i lhs_mat_23_01 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_01), lhs_mat_ymm_23_01, 1);
- __m512i lhs_mat_01_02 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_02), lhs_mat_ymm_01_02, 1);
- __m512i lhs_mat_23_02 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_02), lhs_mat_ymm_23_02, 1);
- __m512i lhs_mat_01_03 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_03), lhs_mat_ymm_01_03, 1);
- __m512i lhs_mat_23_03 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_03), lhs_mat_ymm_23_03, 1);
-
- __m512i lhs_mat_01_10 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_10), lhs_mat_ymm_01_10, 1);
- __m512i lhs_mat_23_10 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_10), lhs_mat_ymm_23_10, 1);
- __m512i lhs_mat_01_11 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_11), lhs_mat_ymm_01_11, 1);
- __m512i lhs_mat_23_11 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_11), lhs_mat_ymm_23_11, 1);
- __m512i lhs_mat_01_12 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_12), lhs_mat_ymm_01_12, 1);
- __m512i lhs_mat_23_12 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_12), lhs_mat_ymm_23_12, 1);
- __m512i lhs_mat_01_13 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_13), lhs_mat_ymm_01_13, 1);
- __m512i lhs_mat_23_13 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_13), lhs_mat_ymm_23_13, 1);
+ __m512i lhs_mat_zmm_0123_00_01 = _mm512_loadu_si512((const __m512i * )((a_ptrs[rp][b].qs + sb * 256)));
+ __m512i lhs_mat_zmm_0123_02_03 = _mm512_loadu_si512((const __m512i * )((a_ptrs[rp][b].qs + 64 + sb * 256)));
+ __m512i lhs_mat_zmm_0123_10_11 = _mm512_loadu_si512((const __m512i * )((a_ptrs[rp][b].qs + 128 + sb * 256)));
+ __m512i lhs_mat_zmm_0123_12_13 = _mm512_loadu_si512((const __m512i * )((a_ptrs[rp][b].qs + 192 + sb * 256)));
+
+ // Loaded as set of 128 bit vectors and repeated and stored into a 512 bit vector
+ __m512i lhs_mat_01_00 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_00_01, lhs_mat_zmm_0123_00_01, 0b00000000);
+ __m512i lhs_mat_23_00 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_00_01, lhs_mat_zmm_0123_00_01, 0b01010101);
+ __m512i lhs_mat_01_01 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_00_01, lhs_mat_zmm_0123_00_01, 0b10101010);
+ __m512i lhs_mat_23_01 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_00_01, lhs_mat_zmm_0123_00_01, 0b11111111);
+
+ __m512i lhs_mat_01_02 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_02_03, lhs_mat_zmm_0123_02_03, 0b00000000);
+ __m512i lhs_mat_23_02 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_02_03, lhs_mat_zmm_0123_02_03, 0b01010101);
+ __m512i lhs_mat_01_03 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_02_03, lhs_mat_zmm_0123_02_03, 0b10101010);
+ __m512i lhs_mat_23_03 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_02_03, lhs_mat_zmm_0123_02_03, 0b11111111);
+
+ __m512i lhs_mat_01_10 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_10_11, lhs_mat_zmm_0123_10_11, 0b00000000);
+ __m512i lhs_mat_23_10 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_10_11, lhs_mat_zmm_0123_10_11, 0b01010101);
+ __m512i lhs_mat_01_11 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_10_11, lhs_mat_zmm_0123_10_11, 0b10101010);
+ __m512i lhs_mat_23_11 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_10_11, lhs_mat_zmm_0123_10_11, 0b11111111);
+
+ __m512i lhs_mat_01_12 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_12_13, lhs_mat_zmm_0123_12_13, 0b00000000);
+ __m512i lhs_mat_23_12 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_12_13, lhs_mat_zmm_0123_12_13, 0b01010101);
+ __m512i lhs_mat_01_13 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_12_13, lhs_mat_zmm_0123_12_13, 0b10101010);
+ __m512i lhs_mat_23_13 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_12_13, lhs_mat_zmm_0123_12_13, 0b11111111);
// Bsums are loaded - four bsums are loaded (for two sub blocks) for the different Q8_K blocks
__m256i lhs_bsums_0123_01 = _mm256_loadu_si256((const __m256i * )((a_ptrs[rp][b].bsums + 16 * sb)));
@@ -2478,111 +2439,85 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
// Loop to iterate over the eight sub blocks of a super block - two sub blocks are processed per iteration
for (int sb = 0; sb < QK_K / 64; sb++) {
+ const __m512i rhs_raw_mat_01234567_0 = _mm512_loadu_si512((const __m512i * )(b_ptr_0[b].qs + sb * 256));
+ const __m512i rhs_raw_mat_01234567_1 = _mm512_loadu_si512((const __m512i * )(b_ptr_0[b].qs + 64 + sb * 256));
+ const __m512i rhs_raw_mat_01234567_2 = _mm512_loadu_si512((const __m512i * )(b_ptr_0[b].qs + 128 + sb * 256));
+ const __m512i rhs_raw_mat_01234567_3 = _mm512_loadu_si512((const __m512i * )(b_ptr_0[b].qs + 192 + sb * 256));
- const __m256i rhs_raw_mat_0123_0 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + sb * 256));
- const __m256i rhs_raw_mat_4567_0 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 32 + sb * 256));
- const __m256i rhs_raw_mat_0123_1 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 64 + sb * 256));
- const __m256i rhs_raw_mat_4567_1 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 96 + sb * 256));
- const __m256i rhs_raw_mat_0123_2 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 128 + sb * 256));
- const __m256i rhs_raw_mat_4567_2 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 160 + sb * 256));
- const __m256i rhs_raw_mat_0123_3 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 192 + sb * 256));
- const __m256i rhs_raw_mat_4567_3 = _mm256_loadu_si256((const __m256i * )(b_ptr_0[b].qs + 224 + sb * 256));
+ const __m512i rhs_raw_mat_89ABCDEF_0 = _mm512_loadu_si512((const __m512i * )(b_ptr_1[b].qs + sb * 256));
+ const __m512i rhs_raw_mat_89ABCDEF_1 = _mm512_loadu_si512((const __m512i * )(b_ptr_1[b].qs + 64 + sb * 256));
+ const __m512i rhs_raw_mat_89ABCDEF_2 = _mm512_loadu_si512((const __m512i * )(b_ptr_1[b].qs + 128 + sb * 256));
+ const __m512i rhs_raw_mat_89ABCDEF_3 = _mm512_loadu_si512((const __m512i * )(b_ptr_1[b].qs + 192 + sb * 256));
- const __m256i rhs_raw_mat_89AB_0 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + sb * 256));
- const __m256i rhs_raw_mat_CDEF_0 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 32 + sb * 256));
- const __m256i rhs_raw_mat_89AB_1 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 64 + sb * 256));
- const __m256i rhs_raw_mat_CDEF_1 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 96 + sb * 256));
- const __m256i rhs_raw_mat_89AB_2 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 128 + sb * 256));
- const __m256i rhs_raw_mat_CDEF_2 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 160 + sb * 256));
- const __m256i rhs_raw_mat_89AB_3 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 192 + sb * 256));
- const __m256i rhs_raw_mat_CDEF_3 = _mm256_loadu_si256((const __m256i * )(b_ptr_1[b].qs + 224 + sb * 256));
+ const __m512i rhs_raw_mat_014589CD_0_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_0, rhsLeftOrderSp1, rhs_raw_mat_89ABCDEF_0);
+ const __m512i rhs_raw_mat_014589CD_0_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_0, rhsLeftOrderSp2, rhs_raw_mat_89ABCDEF_0);
+ const __m512i rhs_raw_mat_2367ABEF_0_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_0, rhsRightOrderSp1, rhs_raw_mat_89ABCDEF_0);
+ const __m512i rhs_raw_mat_2367ABEF_0_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_0, rhsRightOrderSp2, rhs_raw_mat_89ABCDEF_0);
- const __m256i rhs_raw_mat_0145_0 = _mm256_blend_epi32(rhs_raw_mat_0123_0, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_0, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_0 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_0, requiredOrder), rhs_raw_mat_4567_0, 240);
- const __m256i rhs_raw_mat_0145_1 = _mm256_blend_epi32(rhs_raw_mat_0123_1, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_1, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_1 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_1, requiredOrder), rhs_raw_mat_4567_1, 240);
- const __m256i rhs_raw_mat_0145_2 = _mm256_blend_epi32(rhs_raw_mat_0123_2, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_2, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_2 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_2, requiredOrder), rhs_raw_mat_4567_2, 240);
- const __m256i rhs_raw_mat_0145_3 = _mm256_blend_epi32(rhs_raw_mat_0123_3, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_3, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_3 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_3, requiredOrder), rhs_raw_mat_4567_3, 240);
+ const __m512i rhs_raw_mat_014589CD_1_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_1, rhsLeftOrderSp1, rhs_raw_mat_89ABCDEF_1);
+ const __m512i rhs_raw_mat_014589CD_1_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_1, rhsLeftOrderSp2, rhs_raw_mat_89ABCDEF_1);
+ const __m512i rhs_raw_mat_2367ABEF_1_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_1, rhsRightOrderSp1, rhs_raw_mat_89ABCDEF_1);
+ const __m512i rhs_raw_mat_2367ABEF_1_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_1, rhsRightOrderSp2, rhs_raw_mat_89ABCDEF_1);
- const __m256i rhs_raw_mat_89CD_0 = _mm256_blend_epi32(rhs_raw_mat_89AB_0, _mm256_permutevar8x32_epi32(rhs_raw_mat_CDEF_0, requiredOrder), 240);
- const __m256i rhs_raw_mat_ABEF_0 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_89AB_0, requiredOrder), rhs_raw_mat_CDEF_0, 240);
- const __m256i rhs_raw_mat_89CD_1 = _mm256_blend_epi32(rhs_raw_mat_89AB_1, _mm256_permutevar8x32_epi32(rhs_raw_mat_CDEF_1, requiredOrder), 240);
- const __m256i rhs_raw_mat_ABEF_1 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_89AB_1, requiredOrder), rhs_raw_mat_CDEF_1, 240);
- const __m256i rhs_raw_mat_89CD_2 = _mm256_blend_epi32(rhs_raw_mat_89AB_2, _mm256_permutevar8x32_epi32(rhs_raw_mat_CDEF_2, requiredOrder), 240);
- const __m256i rhs_raw_mat_ABEF_2 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_89AB_2, requiredOrder), rhs_raw_mat_CDEF_2, 240);
- const __m256i rhs_raw_mat_89CD_3 = _mm256_blend_epi32(rhs_raw_mat_89AB_3, _mm256_permutevar8x32_epi32(rhs_raw_mat_CDEF_3, requiredOrder), 240);
- const __m256i rhs_raw_mat_ABEF_3 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_89AB_3, requiredOrder), rhs_raw_mat_CDEF_3, 240);
+ const __m512i rhs_raw_mat_014589CD_2_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_2, rhsLeftOrderSp1, rhs_raw_mat_89ABCDEF_2);
+ const __m512i rhs_raw_mat_014589CD_2_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_2, rhsLeftOrderSp2, rhs_raw_mat_89ABCDEF_2);
+ const __m512i rhs_raw_mat_2367ABEF_2_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_2, rhsRightOrderSp1, rhs_raw_mat_89ABCDEF_2);
+ const __m512i rhs_raw_mat_2367ABEF_2_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_2, rhsRightOrderSp2, rhs_raw_mat_89ABCDEF_2);
- const __m512i rhs_raw_mat_014589CD_0 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_0145_0), rhs_raw_mat_89CD_0, 1);
- const __m512i rhs_raw_mat_2367ABEF_0 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_2367_0), rhs_raw_mat_ABEF_0, 1);
- const __m512i rhs_raw_mat_014589CD_1 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_0145_1), rhs_raw_mat_89CD_1, 1);
- const __m512i rhs_raw_mat_2367ABEF_1 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_2367_1), rhs_raw_mat_ABEF_1, 1);
+ const __m512i rhs_raw_mat_014589CD_3_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_3, rhsLeftOrderSp1, rhs_raw_mat_89ABCDEF_3);
+ const __m512i rhs_raw_mat_014589CD_3_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_3, rhsLeftOrderSp2, rhs_raw_mat_89ABCDEF_3);
+ const __m512i rhs_raw_mat_2367ABEF_3_sp1 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_3, rhsRightOrderSp1, rhs_raw_mat_89ABCDEF_3);
+ const __m512i rhs_raw_mat_2367ABEF_3_sp2 = _mm512_permutex2var_epi32(rhs_raw_mat_01234567_3, rhsRightOrderSp2, rhs_raw_mat_89ABCDEF_3);
- const __m512i rhs_raw_mat_014589CD_2 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_0145_2), rhs_raw_mat_89CD_2, 1);
- const __m512i rhs_raw_mat_2367ABEF_2 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_2367_2), rhs_raw_mat_ABEF_2, 1);
- const __m512i rhs_raw_mat_014589CD_3 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_0145_3), rhs_raw_mat_89CD_3, 1);
- const __m512i rhs_raw_mat_2367ABEF_3 = _mm512_inserti32x8(_mm512_castsi256_si512(rhs_raw_mat_2367_3), rhs_raw_mat_ABEF_3, 1);
+ // 4bit-8bit pattern one - right side input
+ const __m512i rhs_mat_014589CD_00_sp1 = _mm512_and_si512(rhs_raw_mat_014589CD_0_sp1, m4bexpanded); //B00(0-3) B01(0-3) B00(0-3) B01(0-3) B04(0-3) B05(0-3) B04(0-3) B05(0-3) B08(0-3) B09(0-3) B08(0-3) B09(0-3) B0C(0-3) B0D(0-3) B0C(0-3) B0D(0-3)
+ const __m512i rhs_mat_014589CD_10_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_0_sp1, 4), m4bexpanded); //B10(0-3) B11(0-3) B10(0-3) B11(0-3) B14(0-3) B15(0-3) B14(0-3) B15(0-3) B18(0-3) B19(0-3) B18(0-3) B19(0-3) B1C(0-3) B1D(0-3) B1C(0-3) B1D(0-3)
- //4-bit -> 8-bit
- const __m512i rhs_mat_014589CD_00 = _mm512_and_si512(rhs_raw_mat_014589CD_0, m4bexpanded); //B00(0-7) B01(0-7) B04(0-7) B05(0-7) B08(0-7) B09(0-7) B0C(0-7) B0D(0-7)
- const __m512i rhs_mat_2367ABEF_00 = _mm512_and_si512(rhs_raw_mat_2367ABEF_0, m4bexpanded); //B02(0-7) B03(0-7) B06(0-7) B07(0-7) B0A(0-7) B0B(0-7) B0E(0-7) B0F(0-7)
- const __m512i rhs_mat_014589CD_01 = _mm512_and_si512(rhs_raw_mat_014589CD_1, m4bexpanded); //B00(8-15) B01(8-15) B04(8-15) B05(8-15) B08(8-15) B09(8-15) B0C(8-15) B0D(8-15)
- const __m512i rhs_mat_2367ABEF_01 = _mm512_and_si512(rhs_raw_mat_2367ABEF_1, m4bexpanded); //B02(8-15) B03(8-15) B06(8-15) B07(8-15) B0A(8-15) B0B(8-15) B0E(8-15) B0F(8-15)
+ const __m512i rhs_mat_2367ABEF_00_sp1 = _mm512_and_si512(rhs_raw_mat_2367ABEF_0_sp1, m4bexpanded); //B02(0-3) B03(0-3) B02(0-3) B03(0-3) B06(0-3) B07(0-3) B06(0-3) B07(0-3) B0A(0-3) B0B(0-3) B0A(0-3) B0B(0-3) B0E(0-3) B0F(0-3) B0E(0-3) B0F(0-3)
+ const __m512i rhs_mat_2367ABEF_10_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_0_sp1, 4), m4bexpanded); //B12(0-3) B13(0-3) B12(0-3) B13(0-3) B16(0-3) B17(0-3) B16(0-3) B17(0-3) B1A(0-3) B1B(0-3) B1A(0-3) B1B(0-3) B1E(0-3) B1F(0-3) B1E(0-3) B1F(0-3)
- const __m512i rhs_mat_014589CD_02 = _mm512_and_si512(rhs_raw_mat_014589CD_2, m4bexpanded); //B00(16-23) B01(16-23) B04(16-23) B05(16-23) B08(16-23) B09(16-23) B0C(16-23) B0D(16-23)
- const __m512i rhs_mat_2367ABEF_02 = _mm512_and_si512(rhs_raw_mat_2367ABEF_2, m4bexpanded); //B02(16-23) B03(16-23) B06(16-23) B07(16-23) B0A(16-23) B0B(16-23) B0E(16-23) B0F(16-23)
- const __m512i rhs_mat_014589CD_03 = _mm512_and_si512(rhs_raw_mat_014589CD_3, m4bexpanded); //B00(24-31) B01(24-31) B04(24-31) B05(24-31) B08(24-31) B09(24-31) B0C(24-31) B0D(24-31)
- const __m512i rhs_mat_2367ABEF_03 = _mm512_and_si512(rhs_raw_mat_2367ABEF_3, m4bexpanded); //B02(24-31) B03(24-31) B06(24-31) B07(24-31) B0A(24-31) B0B(24-31) B0E(24-31) B0F(24-31)
+ const __m512i rhs_mat_014589CD_01_sp1 = _mm512_and_si512(rhs_raw_mat_014589CD_1_sp1, m4bexpanded); //B00(8-11) B01(8-11) B00(8-11) B01(8-11) B04(8-11) B05(8-11) B04(8-11) B05(8-11) B08(8-11) B09(8-11) B08(8-11) B09(8-11) B0C(8-11) B0D(8-11) B0C(8-11) B0D(8-11)
+ const __m512i rhs_mat_014589CD_11_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_1_sp1, 4), m4bexpanded); //B10(8-11) B11(8-11) B10(8-11) B11(8-11) B14(8-11) B15(8-11) B14(8-11) B15(8-11) B18(8-11) B19(8-11) B18(8-11) B19(8-11) B1C(8-11) B1D(8-11) B1C(8-11) B1D(8-11)
- const __m512i rhs_mat_014589CD_10 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_0, 4), m4bexpanded); //B10(0-7) B11(0-7) B14(0-7) B15(0-7) B18(0-7) B19(0-7) B1C(0-7) B1D(0-7)
- const __m512i rhs_mat_2367ABEF_10 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_0, 4), m4bexpanded); //B12(0-7) B13(0-7) B16(0-7) B17(0-7) B1A(0-7) B1B(0-7) B1E(0-7) B1F(0-7)
- const __m512i rhs_mat_014589CD_11 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_1, 4), m4bexpanded); //B10(8-15) B11(8-15) B14(8-15) B15(8-15) B18(8-15) B19(8-15) B1C(8-15) B1D(8-15)
- const __m512i rhs_mat_2367ABEF_11 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_1, 4), m4bexpanded); //B12(8-15) B13(8-15) B16(8-15) B17(8-15) B1A(8-15) B1B(8-15) B1E(8-15) B1F(8-15)
+ const __m512i rhs_mat_2367ABEF_01_sp1 = _mm512_and_si512(rhs_raw_mat_2367ABEF_1_sp1, m4bexpanded); //B02(8-11) B03(8-11) B02(8-11) B03(8-11) B06(8-11) B07(8-11) B06(8-11) B07(8-11) B0A(8-11) B0B(8-11) B0A(8-11) B0B(8-11) B0E(8-11) B0F(8-11) B0E(8-11) B0F(8-11)
+ const __m512i rhs_mat_2367ABEF_11_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_1_sp1, 4), m4bexpanded); //B12(8-11) B13(8-11) B12(8-11) B13(8-11) B16(8-11) B17(8-11) B16(8-11) B17(8-11) B1A(8-11) B1B(8-11) B1A(8-11) B1B(8-1１) B１E(8-１１) B１F(8-１１) B１E(8-１１) B１F(8-１１)
- const __m512i rhs_mat_014589CD_12 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_2, 4), m4bexpanded); //B10(16-23) B11(16-23) B14(16-23) B15(16-23) B18(16-23) B19(16-23) B1C(16-23) B1D(16-23)
- const __m512i rhs_mat_2367ABEF_12 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_2, 4), m4bexpanded); //B12(16-23) B13(16-23) B16(16-23) B17(16-23) B1A(16-23) B1B(16-23) B1E(16-23) B1F(16-23)
- const __m512i rhs_mat_014589CD_13 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_3, 4), m4bexpanded); //B10(24-31) B11(24-31) B14(24-31) B15(24-31) B18(24-31) B19(24-31) B1C(24-31) B1D(24-31)
- const __m512i rhs_mat_2367ABEF_13 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_3, 4), m4bexpanded); //B12(24-31) B13(24-31) B16(24-31) B17(24-31) B1A(24-31) B1B(24-31) B1E(24-31) B1F(24-31)
+ const __m512i rhs_mat_014589CD_02_sp1 = _mm512_and_si512(rhs_raw_mat_014589CD_2_sp1, m4bexpanded); //B00(16-19) B01(16-19) B00(16-19) B01(16-19) B04(16-19) B05(16-19) B04(16-19) B05(16-19) B08(16-19) B09(16-19) B08(16-19) B09(16-19) B0C(16-19) B0D(16-19) B0C(16-19) B0D(16-19)
+ const __m512i rhs_mat_014589CD_12_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_2_sp1, 4), m4bexpanded); //B10(16-19) B11(16-19) B10(16-19) B11(16-19) B14(16-19) B15(16-19) B14(16-19) B15(16-19) B18(16-19) B19(16-19) B18(16-19) B19(16-19) B1C(16-19) B1D(16-19) B1C(16-19) B1D(16-19)
- // Shuffle pattern one - right side input
- const __m512i rhs_mat_014589CD_00_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_00, (_MM_PERM_ENUM)136); //B00(0-3) B01(0-3) B00(0-3) B01(0-3) B04(0-3) B05(0-3) B04(0-3) B05(0-3) B08(0-3) B09(0-3) B08(0-3) B09(0-3) B0C(0-3) B0D(0-3) B0C(0-3) B0D(0-3)
- const __m512i rhs_mat_2367ABEF_00_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_00, (_MM_PERM_ENUM)136); //B02(0-3) B03(0-3) B02(0-3) B03(0-3) B06(0-3) B07(0-3) B06(0-3) B07(0-3) B0A(0-3) B0B(0-3) B0A(0-3) B0B(0-3) B0E(0-3) B0F(0-3) B0E(0-3) B0F(0-3)
- const __m512i rhs_mat_014589CD_01_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_01, (_MM_PERM_ENUM)136); //B00(8-11) B01(8-11) B00(8-11) B01(8-11) B04(8-11) B05(8-11) B04(8-11) B05(8-11) B08(8-11) B09(8-11) B08(8-11) B09(8-11) B0C(8-11) B0D(8-11) B0C(8-11) B0D(8-11)
- const __m512i rhs_mat_2367ABEF_01_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_01, (_MM_PERM_ENUM)136); //B02(8-11) B03(8-11) B02(8-11) B03(8-11) B06(8-11) B07(8-11) B06(8-11) B07(8-11) B0A(8-11) B0B(8-11) B0A(8-11) B0B(8-11) B0E(8-11) B0F(8-11) B0E(8-11) B0F(8-11)
- const __m512i rhs_mat_014589CD_02_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_02, (_MM_PERM_ENUM)136); //B00(16-19) B01(16-19) B00(16-19) B01(16-19) B04(16-19) B05(16-19) B04(16-19) B05(16-19) B08(16-19) B09(16-19) B08(16-19) B09(16-19) B0C(16-19) B0D(16-19) B0C(16-19) B0D(16-19)
- const __m512i rhs_mat_2367ABEF_02_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_02, (_MM_PERM_ENUM)136); //B02(16-19) B03(16-19) B02(16-19) B03(16-19) B06(16-19) B07(16-19) B06(16-19) B07(16-19) B0A(16-19) B0B(16-19) B0A(16-19) B0B(16-19) B0E(16-19) B0F(16-19) B0E(16-19) B0F(16-19)
- const __m512i rhs_mat_014589CD_03_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_03, (_MM_PERM_ENUM)136); //B00(24-27) B01(24-27) B00(24-27) B01(24-27) B04(24-27) B05(24-27) B04(24-27) B05(24-27) B08(24-27) B09(24-27) B08(24-27) B09(24-27) B0C(24-27) B0D(24-27) B0C(24-27) B0D(24-27)
- const __m512i rhs_mat_2367ABEF_03_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_03, (_MM_PERM_ENUM)136); //B02(24-27) B03(24-27) B02(24-27) B03(24-27) B06(24-27) B07(24-27) B06(24-27) B07(24-27) B0A(24-27) B0B(24-27) B0A(24-27) B0B(24-27) B0E(24-27) B0F(24-27) B0E(24-27) B0F(24-27)
+ const __m512i rhs_mat_2367ABEF_02_sp1 = _mm512_and_si512(rhs_raw_mat_2367ABEF_2_sp1, m4bexpanded); //B02(16-19) B03(16-19) B02(16-19) B03(16-19) B06(16-19) B07(16-19) B06(16-19) B07(16-19) B0A(16-19) B0B(16-19) B0A(16-19) B0B(16-19) B0E(16-19) B0F(16-19) B0E(16-19) B0F(16-19)
+ const __m512i rhs_mat_2367ABEF_12_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_2_sp1, 4), m4bexpanded); //B12(16-19) B13(16-19) B12(16-19) B13(16-19) B16(16-19) B17(16-19) B16(16-19) B17(16-19) B1A(16-19) B1B(16-19) B1A(16-19) B1B(16-19) B1E(16-19) B1F(16-19) B1E(16-19) B1F(16-19)
- const __m512i rhs_mat_014589CD_10_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_10, (_MM_PERM_ENUM)136); //B10(0-3) B11(0-3) B10(0-3) B11(0-3) B14(0-3) B15(0-3) B14(0-3) B15(0-3) B18(0-3) B19(0-3) B18(0-3) B19(0-3) B1C(0-3) B1D(0-3) B1C(0-3) B1D(0-3)
- const __m512i rhs_mat_2367ABEF_10_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_10, (_MM_PERM_ENUM)136); //B12(0-3) B13(0-3) B12(0-3) B13(0-3) B16(0-3) B17(0-3) B16(0-3) B17(0-3) B1A(0-3) B1B(0-3) B1A(0-3) B1B(0-3) B1E(0-3) B1F(0-3) B1E(0-3) B1F(0-3)
- const __m512i rhs_mat_014589CD_11_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_11, (_MM_PERM_ENUM)136); //B10(8-11) B11(8-11) B10(8-11) B11(8-11) B14(8-11) B15(8-11) B14(8-11) B15(8-11) B18(8-11) B19(8-11) B18(8-11) B19(8-11) B1C(8-11) B1D(8-11) B1C(8-11) B1D(8-11)
- const __m512i rhs_mat_2367ABEF_11_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_11, (_MM_PERM_ENUM)136); //B12(8-11) B13(8-11) B12(8-11) B13(8-11) B16(8-11) B17(8-11) B16(8-11) B17(8-11) B1A(8-11) B1B(8-11) B1A(8-11) B1B(8-11) B1E(8-11) B1F(8-11) B1E(8-11) B1F(8-11)
- const __m512i rhs_mat_014589CD_12_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_12, (_MM_PERM_ENUM)136); //B10(16-19) B11(16-19) B10(16-19) B11(16-19) B14(16-19) B15(16-19) B14(16-19) B15(16-19) B18(16-19) B19(16-19) B18(16-19) B19(16-19) B1C(16-19) B1D(16-19) B1C(16-19) B1D(16-19)
- const __m512i rhs_mat_2367ABEF_12_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_12, (_MM_PERM_ENUM)136); //B12(16-19) B13(16-19) B12(16-19) B13(16-19) B16(16-19) B17(16-19) B16(16-19) B17(16-19) B1A(16-19) B1B(16-19) B1A(16-19) B1B(16-19) B1E(16-19) B1F(16-19) B1E(16-19) B1F(16-19)
- const __m512i rhs_mat_014589CD_13_sp1 = _mm512_shuffle_epi32(rhs_mat_014589CD_13, (_MM_PERM_ENUM)136); //B10(24-27) B11(24-27) B10(24-27) B11(24-27) B14(24-27) B15(24-27) B14(24-27) B15(24-27) B18(24-27) B19(24-27) B18(24-27) B19(24-27) B1C(24-27) B1D(24-27) B1C(24-27) B1D(24-27)
- const __m512i rhs_mat_2367ABEF_13_sp1 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_13, (_MM_PERM_ENUM)136); //B12(24-27) B13(24-27) B12(24-27) B13(24-27) B16(24-27) B17(24-27) B16(24-27) B17(24-27) B1A(24-27) B1B(24-27) B1A(24-27) B1B(24-27) B1E(24-27) B1F(24-27) B1E(24-27) B1F(24-27)
+ const __m512i rhs_mat_014589CD_03_sp1 = _mm512_and_si512(rhs_raw_mat_014589CD_3_sp1, m4bexpanded); //B00(24-27) B01(24-27) B00(24-27) B01(24-27) B04(24-27) B05(24-27) B04(24-27) B05(24-27) B08(24-27) B09(24-27) B08(24-27) B09(24-27) B0C(24-27) B0D(24-27) B0C(24-27) B0D(24-27)
+ const __m512i rhs_mat_014589CD_13_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_3_sp1, 4), m4bexpanded); //B10(24-27) B11(24-27) B10(24-27) B11(24-27) B14(24-27) B15(24-27) B14(24-27) B15(24-27) B18(24-27) B19(24-27) B18(24-27) B19(24-27) B1C(24-27) B1D(24-27) B1C(24-27) B1D(24-27)
- // Shuffle pattern two - right side input
- const __m512i rhs_mat_014589CD_00_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_00, (_MM_PERM_ENUM)221); //B00(4-7) B01(4-7) B00(4-7) B01(4-7) B04(4-7) B05(4-7) B04(4-7) B05(4-7) B08(4-7) B09(4-7) B08(4-7) B09(4-7) B0C(4-7) B0D(4-7) B0C(4-7) B0D(4-7)
- const __m512i rhs_mat_2367ABEF_00_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_00, (_MM_PERM_ENUM)221); //B02(4-7) B03(4-7) B02(4-7) B03(4-7) B06(4-7) B07(4-7) B06(4-7) B07(4-7) B0A(4-7) B0B(4-7) B0A(4-7) B0B(4-7) B0E(4-7) B0F(4-7) B0E(4-7) B0F(4-7)
- const __m512i rhs_mat_014589CD_01_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_01, (_MM_PERM_ENUM)221); //B00(12-15) B01(12-15) B00(12-15) B01(12-15) B04(12-15) B05(12-15) B04(12-15) B05(12-15) B08(12-15) B09(12-15) B08(12-15) B09(12-15) B0C(12-15) B0D(12-15) B0C(12-15) B0D(12-15)
- const __m512i rhs_mat_2367ABEF_01_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_01, (_MM_PERM_ENUM)221); //B02(12-15) B03(12-15) B02(12-15) B03(12-15) B06(12-15) B07(12-15) B06(12-15) B07(12-15) B0A(12-15) B0B(12-15) B0A(12-15) B0B(12-15) B0E(12-15) B0F(12-15) B0E(12-15) B0F(12-15)
- const __m512i rhs_mat_014589CD_02_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_02, (_MM_PERM_ENUM)221); //B00(20-23) B01(20-23) B00(20-23) B01(20-23) B04(20-23) B05(20-23) B04(20-23) B05(20-23) B08(20-23) B09(20-23) B08(20-23) B09(20-23) B0C(20-23) B0D(20-23) B0C(20-23) B0D(20-23)
- const __m512i rhs_mat_2367ABEF_02_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_02, (_MM_PERM_ENUM)221); //B02(20-23) B03(20-23) B02(20-23) B03(20-23) B06(20-23) B07(20-23) B06(20-23) B07(20-23) B0A(20-23) B0B(20-23) B0A(20-23) B0B(20-23) B0E(20-23) B0F(20-23) B0E(20-23) B0F(20-23)
- const __m512i rhs_mat_014589CD_03_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_03, (_MM_PERM_ENUM)221); //B00(28-31) B01(28-31) B00(28-31) B01(28-31) B04(28-31) B05(28-31) B04(28-31) B05(28-31) B08(28-31) B09(28-31) B08(28-31) B09(28-31) B0C(28-31) B0D(28-31) B0C(28-31) 0BD(28-31)
- const __m512i rhs_mat_2367ABEF_03_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_03, (_MM_PERM_ENUM)221); //B02(28-31) B03(28-31) B02(28-31) B03(28-31) B06(28-31) B07(28-31) B06(28-31) B07(28-31) B0A(28-31) B0B(28-31) B0A(28-31) B0B(28-31) B0E(28-31) B0F(28-31) B0E(28-31) B0F(28-31)
+ const __m512i rhs_mat_2367ABEF_03_sp1 = _mm512_and_si512(rhs_raw_mat_2367ABEF_3_sp1, m4bexpanded); //B02(24-27) B03(24-27) B02(24-27) B03(24-27) B06(24-27) B07(24-27) B06(24-27) B07(24-27) B0A(24-27) B0B(24-27) B0A(24-27) B0B(24-27) B0E(24-27) B0F(24-27) B0E(24-27) B0F(24-27)
+ const __m512i rhs_mat_2367ABEF_13_sp1 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_3_sp1, 4), m4bexpanded); //B12(24-27) B13(24-27) B12(24-27) B13(24-27) B16(24-27) B17(24-27) B16(24-27) B17(24-27) B1A(24-27) B1B(24-27) B1A(24-27) B1B(24-27) B1E(24-27) B1F(24-27) B1E(24-27) B1F(24-27)
- const __m512i rhs_mat_014589CD_10_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_10, (_MM_PERM_ENUM)221); //B10(4-7) B11(4-7) B10(4-7) B11(4-7) B14(4-7) B15(4-7) B14(4-7) B15(4-7) B18(4-7) B19(4-7) B18(4-7) B19(4-7) B1C(4-7) B1D(4-7) B1C(4-7) B1D(4-7)
- const __m512i rhs_mat_2367ABEF_10_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_10, (_MM_PERM_ENUM)221); //B12(4-7) B13(4-7) B12(4-7) B13(4-7) B16(4-7) B17(4-7) B16(4-7) B17(4-7) B1A(4-7) B1B(4-7) B1A(4-7) B1B(4-7) B1E(4-7) B1F(4-7) B1E(4-7) B1F(4-7)
- const __m512i rhs_mat_014589CD_11_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_11, (_MM_PERM_ENUM)221); //B10(12-15) B11(12-15) B10(12-15) B11(12-15) B14(12-15) B15(12-15) B14(12-15) B15(12-15) B18(12-15) B19(12-15) B18(12-15) B19(12-15) B1C(12-15) B1D(12-15) B1C(12-15) B1D(12-15)
- const __m512i rhs_mat_2367ABEF_11_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_11, (_MM_PERM_ENUM)221); //B12(12-15) B13(12-15) B12(12-15) B13(12-15) B16(12-15) B17(12-15) B16(12-15) B17(12-15) B1A(12-15) B1B(12-15) B1A(12-15) B1B(12-15) B1E(12-15) B1F(12-15) B1E(12-15) B1F(12-15)
- const __m512i rhs_mat_014589CD_12_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_12, (_MM_PERM_ENUM)221); //B10(20-23) B11(20-23) B10(20-23) B11(20-23) B14(20-23) B15(20-23) B14(20-23) B15(20-23) B18(20-23) B19(20-23) B18(20-23) B19(20-23) B1C(20-23) B1D(20-23) B1C(20-23) B1D(20-23)
- const __m512i rhs_mat_2367ABEF_12_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_12, (_MM_PERM_ENUM)221); //B12(20-23) B13(20-23) B12(20-23) B13(20-23) B16(20-23) B17(20-23) B16(20-23) B17(20-23) B1A(20-23) B1B(20-23) B1A(20-23) B1B(20-23) B1E(20-23) B1F(20-23) B1E(20-23) B1F(20-23)
- const __m512i rhs_mat_014589CD_13_sp2 = _mm512_shuffle_epi32(rhs_mat_014589CD_13, (_MM_PERM_ENUM)221); //B10(28-31) B11(28-31) B10(28-31) B11(28-31) B14(28-31) B15(28-31) B14(28-31) B15(28-31) B18(28-31) B19(28-31) B18(28-31) B19(28-31) B1C(28-31) B1D(28-31) B1C(28-31) B1D(28-31)
- const __m512i rhs_mat_2367ABEF_13_sp2 = _mm512_shuffle_epi32(rhs_mat_2367ABEF_13, (_MM_PERM_ENUM)221); //B12(28-31) B13(28-31) B12(28-31) B13(28-31) B16(28-31) B17(28-31) B16(28-31) B17(28-31) B1A(28-31) B1B(28-31) B1A(28-31) B1B(28-31) B1E(28-31) B1F(28-31) B1E(28-31) B1F(28-31)
+ // 4bit-8bit pattern two - right side input
+ const __m512i rhs_mat_014589CD_00_sp2 = _mm512_and_si512(rhs_raw_mat_014589CD_0_sp2, m4bexpanded); //B00(4-7) B01(4-7) B00(4-7) B01(4-7) B04(4-7) B05(4-7) B04(4-7) B05(4-7) B08(4-7) B09(4-7) B08(4-7) B09(4-7) B0C(4-7) B0D(4-7) B0C(4-7) B0D(4-7)
+ const __m512i rhs_mat_014589CD_10_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_0_sp2, 4), m4bexpanded); //B10(4-7) B11(4-7) B10(4-7) B11(4-7) B14(4-7) B15(4-7) B14(4-7) B15(4-7) B18(4-7) B19(4-7) B18(4-7) B19(4-7) B1C(4-7) B1D(4-7) B1C(4-7) B1D(4-7)
+
+ const __m512i rhs_mat_2367ABEF_00_sp2 = _mm512_and_si512(rhs_raw_mat_2367ABEF_0_sp2, m4bexpanded); //B02(4-7) B03(4-7) B02(4-7) B03(4-7) B06(4-7) B07(4-7) B06(4-7) B07(4-7) B0A(4-7) B0B(4-7) B0A(4-7) B0B(4-7) B0E(4-7) B0F(4-7) B0E(4-7] B0F(4-7)
+ const __m512i rhs_mat_2367ABEF_10_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_0_sp2, 4), m4bexpanded); //B12(4-7] B13(4-7] B12(4-7] B13(4-7] B16(4-7] B17(4-七年] B16(四年] B1七七年] B１A（四年]B１B（四年]B１A（四年]B１B（四年]B１E（四年]B１F（四年]B１E（四年]B１F（四年]
+
+ const __m512i rhs_mat_014589CD_01_sp2 = _mm512_and_si512(rhs_raw_mat_014589CD_1_sp2, m4bexpanded); //B00(12-15) B01(12-15) B00(12-15) B01(12-15) B04(12-15) B05(12-15) B04(12-15) B05(12-15) B08(12-15) B09(12-15) B08(12-15) B09(12-15) B0C(12-15) B0D(12-15) B0C(12-15) B0D(12-15)
+ const __m512i rhs_mat_014589CD_11_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_1_sp2, 4), m4bexpanded); //B10(12-15) B11(12-15) B10(12-15) B11(12-15) B14(12-15) B15(12-15) B14(12-15) B15(12-15) B18(12-15) B19(12-15) B18(12-15) B19(12-15) B1C(12-15) B1D(12-15) B1C(12-15) B1D(12-15)
+
+ const __m512i rhs_mat_2367ABEF_01_sp2 = _mm512_and_si512(rhs_raw_mat_2367ABEF_1_sp2, m4bexpanded); //B02(12-15) B03(12-15) B02(12-15) B03(12-15) B06(12-15) B07(12-15) B06(12-15) B07(12-15) B0A(12-15) B0B(12-15) B0A(12-15) B0B(12-15) B0E(12-15) B0F(12-15) B0E(12-15) B0F(12-15)
+ const __m512i rhs_mat_2367ABEF_11_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_1_sp2, 4), m4bexpanded); //B12(12-15) B13(12-15) B12(12-15) B13(12-15) B16(12-15) B17(12-15) B16(12-15) B17(12-15) B1A(12-15) B1B(12-15) B1A(12-15) B1B(12-15) B1E(12-15) B1F(12-15) B1E(12-15) B1F(12-15)
+
+ const __m512i rhs_mat_014589CD_02_sp2 = _mm512_and_si512(rhs_raw_mat_014589CD_2_sp2, m4bexpanded); //B00(20-23) B01(20-23) B00(20-23) B01(20-23) B04(20-23) B05(20-23) B04(20-23) B05(20-23) B08(20-23) B09(20-23) B08(20-23) B09(20-23) B0C(20-23) B0D(20-23) B0C(20-23) B0D(20-23)
+ const __m512i rhs_mat_014589CD_12_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_2_sp2, 4), m4bexpanded); //B10(20-23) B11(20-23) B10(20-23) B11(20-23) B14(20-23) B15(20-23) B14(20-23) B15(20-23) B18(20-23) B19(20-23) B18(20-23) B19(20-23) B1C(20-23) B1D(20-23) B1C(20-23) B1D(20-23)
+
+ const __m512i rhs_mat_2367ABEF_02_sp2 = _mm512_and_si512(rhs_raw_mat_2367ABEF_2_sp2, m4bexpanded); //B02(20-23) B03(20-23) B02(20-23) B03(20-23) B06(20-23) B07(20-23) B06(20-23) B07(20-23) B0A(20-23) B0B(20-23) B0A(20-23) B0B(20-23) B0E(20-23) B0F(20-23) B0E(20-23) B0F(20-23)
+ const __m512i rhs_mat_2367ABEF_12_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_2_sp2, 4), m4bexpanded); //B12(20-23) B13(20-23) B12(20-23) B13(20-23) B16(20-23) B17(20-23) B16(20-23) B17(20-23) B1A(20-23) B1B(20-23) B1A(20-23) B1B(20-23) B1E(20-23) B1F(20-23) B1E(20-23) B1F(20-23)
+
+ const __m512i rhs_mat_014589CD_03_sp2 = _mm512_and_si512(rhs_raw_mat_014589CD_3_sp2, m4bexpanded); //B00(28-31) B01(28-31) B00(28-31) B01(28-31) B04(28-31) B05(28-31) B04(28-31) B05(28-31) B08(28-31) B09(28-31) B08(28-31) B09(28-31) B0C(28-31) B0D(28-31) B0C(28-31) 0BD(28-31)
+ const __m512i rhs_mat_014589CD_13_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_014589CD_3_sp2, 4), m4bexpanded); //B10(28-31) B11(28-31) B10(28-31) B11(28-31) B14(28-31) B15(28-31) B14(28-31) B15(28-31) B18(28-31) B19(28-31) B18(28-31) B19(28-31) B1C(28-31) B1D(28-31) B1C(28-31) B1D(28-31)
+
+ const __m512i rhs_mat_2367ABEF_03_sp2 = _mm512_and_si512(rhs_raw_mat_2367ABEF_3_sp2, m4bexpanded); //B02(28-31) B03(28-31) B02(28-31) B03(28-31) B06(28-31) B07(28-31) B06(28-31) B07(28-31) B0A(28-31) B0B(28-31) B0A(28-31) B0B(28-31) B0E(28-31) B0F(28-31) B0E(28-31) B0F(28-31)
+ const __m512i rhs_mat_2367ABEF_13_sp2 = _mm512_and_si512(_mm512_srli_epi16(rhs_raw_mat_2367ABEF_3_sp2, 4), m4bexpanded); //B12(28-31) B13(28-31) B12(28-31) B13(28-31) B16(28-31) B17(28-31) B16(28-31) B17(28-31) B1A(28-31) B1B(28-31) B1A(28-31) B1B(28-31) B1E(28-31) B1F(28-31) B1E(28-31) B1F(28-31)
uint32_t utmp_00[4], utmp_01[4], utmp_10[4], utmp_11[4];
@@ -2637,50 +2572,31 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
const __m512i scale_2367ABEF_1 = _mm512_shuffle_epi32(scales_1, (_MM_PERM_ENUM)238);
// Load the four block_q8_k quantized values interleaved with each other in chunks of eight bytes - A0,A1,A2,A3
- // Loaded as set of 128 bit vectors and repeated into a 256 bit vector
- __m256i lhs_mat_ymm_0123_00 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].qs + 256 * sb)));
- __m256i lhs_mat_ymm_01_00 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_00, lhs_mat_ymm_0123_00, 0);
- __m256i lhs_mat_ymm_23_00 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_00, lhs_mat_ymm_0123_00, 17);
- __m256i lhs_mat_ymm_0123_01 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].qs + 32 + 256 * sb)));
- __m256i lhs_mat_ymm_01_01 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_01, lhs_mat_ymm_0123_01, 0);
- __m256i lhs_mat_ymm_23_01 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_01, lhs_mat_ymm_0123_01, 17);
- __m256i lhs_mat_ymm_0123_02 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].qs + 64 + 256 * sb)));
- __m256i lhs_mat_ymm_01_02 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_02, lhs_mat_ymm_0123_02, 0);
- __m256i lhs_mat_ymm_23_02 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_02, lhs_mat_ymm_0123_02, 17);
- __m256i lhs_mat_ymm_0123_03 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].qs + 96 + 256 * sb)));
- __m256i lhs_mat_ymm_01_03 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_03, lhs_mat_ymm_0123_03, 0);
- __m256i lhs_mat_ymm_23_03 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_03, lhs_mat_ymm_0123_03, 17);
- __m256i lhs_mat_ymm_0123_10 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].qs + 128 + 256 * sb)));
- __m256i lhs_mat_ymm_01_10 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_10, lhs_mat_ymm_0123_10, 0);
- __m256i lhs_mat_ymm_23_10 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_10, lhs_mat_ymm_0123_10, 17);
- __m256i lhs_mat_ymm_0123_11 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].qs + 160 + 256 * sb)));
- __m256i lhs_mat_ymm_01_11 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_11, lhs_mat_ymm_0123_11, 0);
- __m256i lhs_mat_ymm_23_11 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_11, lhs_mat_ymm_0123_11, 17);
- __m256i lhs_mat_ymm_0123_12 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].qs + 192 + 256 * sb)));
- __m256i lhs_mat_ymm_01_12 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_12, lhs_mat_ymm_0123_12, 0);
- __m256i lhs_mat_ymm_23_12 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_12, lhs_mat_ymm_0123_12, 17);
- __m256i lhs_mat_ymm_0123_13 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].qs + 224 + 256 * sb)));
- __m256i lhs_mat_ymm_01_13 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_13, lhs_mat_ymm_0123_13, 0);
- __m256i lhs_mat_ymm_23_13 = _mm256_permute2f128_si256(lhs_mat_ymm_0123_13, lhs_mat_ymm_0123_13, 17);
-
- //Loaded as set of 128 bit vectors and repeated and stored into a 256 bit vector before again repeating into a 512 bit vector
- __m512i lhs_mat_01_00 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_00), lhs_mat_ymm_01_00, 1);
- __m512i lhs_mat_23_00 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_00), lhs_mat_ymm_23_00, 1);
- __m512i lhs_mat_01_01 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_01), lhs_mat_ymm_01_01, 1);
- __m512i lhs_mat_23_01 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_01), lhs_mat_ymm_23_01, 1);
- __m512i lhs_mat_01_02 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_02), lhs_mat_ymm_01_02, 1);
- __m512i lhs_mat_23_02 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_02), lhs_mat_ymm_23_02, 1);
- __m512i lhs_mat_01_03 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_03), lhs_mat_ymm_01_03, 1);
- __m512i lhs_mat_23_03 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_03), lhs_mat_ymm_23_03, 1);
-
- __m512i lhs_mat_01_10 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_10), lhs_mat_ymm_01_10, 1);
- __m512i lhs_mat_23_10 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_10), lhs_mat_ymm_23_10, 1);
- __m512i lhs_mat_01_11 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_11), lhs_mat_ymm_01_11, 1);
- __m512i lhs_mat_23_11 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_11), lhs_mat_ymm_23_11, 1);
- __m512i lhs_mat_01_12 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_12), lhs_mat_ymm_01_12, 1);
- __m512i lhs_mat_23_12 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_12), lhs_mat_ymm_23_12, 1);
- __m512i lhs_mat_01_13 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_01_13), lhs_mat_ymm_01_13, 1);
- __m512i lhs_mat_23_13 = _mm512_inserti32x8(_mm512_castsi256_si512(lhs_mat_ymm_23_13), lhs_mat_ymm_23_13, 1);
+ __m512i lhs_mat_zmm_0123_00_01 = _mm512_loadu_si512((const __m512i * )((a_ptr[b].qs + sb * 256)));
+ __m512i lhs_mat_zmm_0123_02_03 = _mm512_loadu_si512((const __m512i * )((a_ptr[b].qs + 64 + sb * 256)));
+ __m512i lhs_mat_zmm_0123_10_11 = _mm512_loadu_si512((const __m512i * )((a_ptr[b].qs + 128 + sb * 256)));
+ __m512i lhs_mat_zmm_0123_12_13 = _mm512_loadu_si512((const __m512i * )((a_ptr[b].qs + 192 + sb * 256)));
+
+ // Loaded as set of 128 bit vectors and repeated and stored into a 512 bit vector
+ __m512i lhs_mat_01_00 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_00_01, lhs_mat_zmm_0123_00_01, 0b00000000);
+ __m512i lhs_mat_23_00 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_00_01, lhs_mat_zmm_0123_00_01, 0b01010101);
+ __m512i lhs_mat_01_01 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_00_01, lhs_mat_zmm_0123_00_01, 0b10101010);
+ __m512i lhs_mat_23_01 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_00_01, lhs_mat_zmm_0123_00_01, 0b11111111);
+
+ __m512i lhs_mat_01_02 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_02_03, lhs_mat_zmm_0123_02_03, 0b00000000);
+ __m512i lhs_mat_23_02 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_02_03, lhs_mat_zmm_0123_02_03, 0b01010101);
+ __m512i lhs_mat_01_03 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_02_03, lhs_mat_zmm_0123_02_03, 0b10101010);
+ __m512i lhs_mat_23_03 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_02_03, lhs_mat_zmm_0123_02_03, 0b11111111);
+
+ __m512i lhs_mat_01_10 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_10_11, lhs_mat_zmm_0123_10_11, 0b00000000);
+ __m512i lhs_mat_23_10 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_10_11, lhs_mat_zmm_0123_10_11, 0b01010101);
+ __m512i lhs_mat_01_11 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_10_11, lhs_mat_zmm_0123_10_11, 0b10101010);
+ __m512i lhs_mat_23_11 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_10_11, lhs_mat_zmm_0123_10_11, 0b11111111);
+
+ __m512i lhs_mat_01_12 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_12_13, lhs_mat_zmm_0123_12_13, 0b00000000);
+ __m512i lhs_mat_23_12 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_12_13, lhs_mat_zmm_0123_12_13, 0b01010101);
+ __m512i lhs_mat_01_13 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_12_13, lhs_mat_zmm_0123_12_13, 0b10101010);
+ __m512i lhs_mat_23_13 = _mm512_shuffle_i32x4(lhs_mat_zmm_0123_12_13, lhs_mat_zmm_0123_12_13, 0b11111111);
// Bsums are loaded - four bsums are loaded (for two sub blocks) for the different Q8_K blocks
__m256i lhs_bsums_0123_01 = _mm256_loadu_si256((const __m256i * )((a_ptr[b].bsums + 16 * sb)));
@@ -2812,7 +2728,7 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
xstart = anc/8;
y = 0;
}
-#endif // __AVX512BW__ && __AVX512DQ__
+#endif // __AVX512F__ && __AVX512BW__ && __AVX512DQ__
// Take group of four block_q8_Kx4 structures at each pass of the loop and perform dot product operation
for (; y < anr / 4; y += 4) {
@@ -2863,14 +2779,14 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
const __m256i rhs_raw_mat_4567_3 = _mm256_loadu_si256((const __m256i * )(b_ptr[b].qs + 224 + sb * 256));
// Save the values in the following vectors in the formats B0B1B4B5, B2B3B6B7 for further processing and storing of values
- const __m256i rhs_raw_mat_0145_0 = _mm256_blend_epi32(rhs_raw_mat_0123_0, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_0, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_0 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_0, requiredOrder), rhs_raw_mat_4567_0, 240);
- const __m256i rhs_raw_mat_0145_1 = _mm256_blend_epi32(rhs_raw_mat_0123_1, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_1, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_1 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_1, requiredOrder), rhs_raw_mat_4567_1, 240);
- const __m256i rhs_raw_mat_0145_2 = _mm256_blend_epi32(rhs_raw_mat_0123_2, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_2, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_2 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_2, requiredOrder), rhs_raw_mat_4567_2, 240);
- const __m256i rhs_raw_mat_0145_3 = _mm256_blend_epi32(rhs_raw_mat_0123_3, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_3, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_3 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_3, requiredOrder), rhs_raw_mat_4567_3, 240);
+ const __m256i rhs_raw_mat_0145_0 = _mm256_permute2x128_si256(rhs_raw_mat_0123_0, rhs_raw_mat_4567_0, 0x20); //B00(0-7) B01(0-7) B04(0-7) B05(0-7)
+ const __m256i rhs_raw_mat_2367_0 = _mm256_permute2x128_si256(rhs_raw_mat_0123_0, rhs_raw_mat_4567_0, 0x31); //B02(0-7) B03(0-7) B06(0-7) B07(0-7)
+ const __m256i rhs_raw_mat_0145_1 = _mm256_permute2x128_si256(rhs_raw_mat_0123_1, rhs_raw_mat_4567_1, 0x20); //B00(8-15) B01(8-15) B04(8-15) B05(8-15)
+ const __m256i rhs_raw_mat_2367_1 = _mm256_permute2x128_si256(rhs_raw_mat_0123_1, rhs_raw_mat_4567_1, 0x31); //B02(8-15) B03(8-15) B06(8-15) B07(8-15)
+ const __m256i rhs_raw_mat_0145_2 = _mm256_permute2x128_si256(rhs_raw_mat_0123_2, rhs_raw_mat_4567_2, 0x20); //B00(16-23) B01(16-23) B04(16-23) B05(16-23)
+ const __m256i rhs_raw_mat_2367_2 = _mm256_permute2x128_si256(rhs_raw_mat_0123_2, rhs_raw_mat_4567_2, 0x31); //B02(16-23) B03(16-23) B06(16-23) B07(16-23)
+ const __m256i rhs_raw_mat_0145_3 = _mm256_permute2x128_si256(rhs_raw_mat_0123_3, rhs_raw_mat_4567_3, 0x20); //B00(24-31) B01(24-31) B04(24-31) B05(24-31)
+ const __m256i rhs_raw_mat_2367_3 = _mm256_permute2x128_si256(rhs_raw_mat_0123_3, rhs_raw_mat_4567_3, 0x31); //B02(24-31) B03(24-31) B06(24-31) B07(24-31)
// 4-bit -> 8-bit
// First sub block of the two sub blocks processed in the iteration
@@ -3031,10 +2947,10 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
const __m256i lhs_mat_23_02_sp1 = _mm256_shuffle_epi32(lhs_mat_23_02, 160); //A02(16-19) A03(16-19) A02(16-19) A03(16-19) A02(16-19) A03(16-19) A02(16-19) A03(16-19)
const __m256i lhs_mat_01_03_sp1 = _mm256_shuffle_epi32(lhs_mat_01_03, 160); //A00(24-27) A00(24-27) A01(24-27) A01(24-27) A00(24-27) A00(24-27) A01(24-27) A01(24-27)
- const __m256i lhs_mat_23_03_sp1 = _mm256_shuffle_epi32(lhs_mat_23_03, 160); //A02(24-27) A03(24-27) A02(24-27) A03(24-27) A02(24-27) A03(24-27) A02(24-27) A03(24-27)
+ const __m256i lhs_mat_23_03_sp1 = _mm256_shuffle_epi32(lhs_mat_23_03, 160); //A02(24-27) A03(24-27) A02(24-27) A03(24-27) A02(24-27) A03(24-27) A02(24-27) A03(24-27)
- const __m256i lhs_mat_01_10_sp1 = _mm256_shuffle_epi32(lhs_mat_01_10, 160); //A10(0-3) A10(0-3) A11(0-3) A11(0-3) A10(0-3) A10(0-3) A11(0-3) A11(0-3)
- const __m256i lhs_mat_23_10_sp1 = _mm256_shuffle_epi32(lhs_mat_23_10, 160); //A12(0-3) A13(0-3) A12(0-3) A13(0-3) A12(0-3) A13(0-3) A12(0-3) A13(0-3)
+ const __m256i lhs_mat_01_10_sp1 = _mm256_shuffle_epi32(lhs_mat_01_10, 160); //A10(0-3) A10(0-3) A11(0-3) A11(0-3) A10(0-3) A10(0-3) A11(0-3) A11(0-3)
+ const __m256i lhs_mat_23_10_sp1 = _mm256_shuffle_epi32(lhs_mat_23_10, 160); //A12(0-3) A13(0-3) A12(0-3) A13(0-3) A12(0-3) A13(0-3) A12(0-3) A13(0-3)
const __m256i lhs_mat_01_11_sp1 = _mm256_shuffle_epi32(lhs_mat_01_11, 160); //A10(8-11) A10(8-11) A11(8-11) A11(8-11) A10(8-11) A10(8-11) A11(8-11) A11(8-11)
const __m256i lhs_mat_23_11_sp1 = _mm256_shuffle_epi32(lhs_mat_23_11, 160); //A12(8-11) A13(8-11) A12(8-11) A13(8-11) A12(8-11) A13(8-11) A12(8-11) A13(8-11)
@@ -3155,6 +3071,7 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
}
}
}
+
for (; y < nr / 4; y++) {
const block_q8_Kx4 * a_ptr = a_ptr_start + (y * nb);
@@ -3196,14 +3113,14 @@ void ggml_gemm_q4_K_8x8_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const vo
const __m256i rhs_raw_mat_4567_3 = _mm256_loadu_si256((const __m256i * )(b_ptr[b].qs + 224 + sb * 256));
// Save the values in the following vectors in the formats B0B1B4B5, B2B3B6B7 for further processing and storing of values
- const __m256i rhs_raw_mat_0145_0 = _mm256_blend_epi32(rhs_raw_mat_0123_0, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_0, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_0 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_0, requiredOrder), rhs_raw_mat_4567_0, 240);
- const __m256i rhs_raw_mat_0145_1 = _mm256_blend_epi32(rhs_raw_mat_0123_1, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_1, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_1 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_1, requiredOrder), rhs_raw_mat_4567_1, 240);
- const __m256i rhs_raw_mat_0145_2 = _mm256_blend_epi32(rhs_raw_mat_0123_2, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_2, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_2 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_2, requiredOrder), rhs_raw_mat_4567_2, 240);
- const __m256i rhs_raw_mat_0145_3 = _mm256_blend_epi32(rhs_raw_mat_0123_3, _mm256_permutevar8x32_epi32(rhs_raw_mat_4567_3, requiredOrder), 240);
- const __m256i rhs_raw_mat_2367_3 = _mm256_blend_epi32(_mm256_permutevar8x32_epi32(rhs_raw_mat_0123_3, requiredOrder), rhs_raw_mat_4567_3, 240);
+ const __m256i rhs_raw_mat_0145_0 = _mm256_permute2x128_si256(rhs_raw_mat_0123_0, rhs_raw_mat_4567_0, 0x20); //B00(0-7) B01(0-7) B04(0-7) B05(0-7)
+ const __m256i rhs_raw_mat_2367_0 = _mm256_permute2x128_si256(rhs_raw_mat_0123_0, rhs_raw_mat_4567_0, 0x31); //B02(0-7) B03(0-7) B06(0-7) B07(0-7)
+ const __m256i rhs_raw_mat_0145_1 = _mm256_permute2x128_si256(rhs_raw_mat_0123_1, rhs_raw_mat_4567_1, 0x20); //B00(8-15) B01(8-15) B04(8-15) B05(8-15)
+ const __m256i rhs_raw_mat_2367_1 = _mm256_permute2x128_si256(rhs_raw_mat_0123_1, rhs_raw_mat_4567_1, 0x31); //B02(8-15) B03(8-15) B06(8-15) B07(8-15)
+ const __m256i rhs_raw_mat_0145_2 = _mm256_permute2x128_si256(rhs_raw_mat_0123_2, rhs_raw_mat_4567_2, 0x20); //B00(16-23) B01(16-23) B04(16-23) B05(16-23)
+ const __m256i rhs_raw_mat_2367_2 = _mm256_permute2x128_si256(rhs_raw_mat_0123_2, rhs_raw_mat_4567_2, 0x31); //B02(16-23) B03(16-23) B06(16-23) B07(16-23)
+ const __m256i rhs_raw_mat_0145_3 = _mm256_permute2x128_si256(rhs_raw_mat_0123_3, rhs_raw_mat_4567_3, 0x20); //B00(24-31) B01(24-31) B04(24-31) B05(24-31)
+ const __m256i rhs_raw_mat_2367_3 = _mm256_permute2x128_si256(rhs_raw_mat_0123_3, rhs_raw_mat_4567_3, 0x31); //B02(24-31) B03(24-31) B06(24-31) B07(24-31)
// 4-bit -> 8-bit
// First sub block of the two sub blocks processed in the iteration
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 <algorithm>
#include <stdint.h>
// 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<int64_t>(n_rows, INT32_MAX));
auto launch_kernel = [&](auto dim) {
concat_non_cont<T, dim><<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
(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 <typename T_IK, typename T_MASK>
+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 <typename T_MASK>
+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<half><<<grid, DSA_TILE_THREADS, smem, ctx.stream()>>>(
+ (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<float><<<grid, DSA_TILE_THREADS, smem, ctx.stream()>>>(
+ (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<T_IK, T_MASK><<<grid, CUDA_DSA_SCORE_BLOCK_SIZE, smem, ctx.stream()>>>(
+ (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 <algorithm>
@@ -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<<<num_blocks, CUDA_SINKHORN_BLOCK_SIZE, 0, ctx.stream()>>>(
+ (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<<<num_blocks, CUDA_SINKHORN_BLOCK_SIZE, 0, ctx.stream()>>>(
+ (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<<<nrows, CUDA_TOP_K_RADIX_BLOCK_SIZE, 0, stream>>>(
+ 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, const char *> 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, const char *> 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, const char *> 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, llm_tensor_info> 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<uint32_t>(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<uint32_t>(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<ggml_tensor *> 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<llm_graph_context> 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<uint32_t, LLAMA_MAX_LAYERS> 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<const llama_memory_hybrid_iswa_context *>(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<int64_t>(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<llm_graph_context> 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<graph_mtp>(*this, params);
+ }
+ return std::make_unique<graph>(*this, params);
+}
+
+llama_model_openpangu_v2::graph::graph(const llama_model & model, const llm_graph_params & params) :
+ graph_base(params, static_cast<const llama_model_openpangu_v2 &>(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<const llama_model_openpangu_v2 &>(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<llm_graph_input_embd>(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<int64_t, 4> ne;
+ const int n_iter;
+ const float eps;
+
+ std::string vars() override { return VARS_TO_STR3(ne, n_iter, eps); }
+
+ test_sinkhorn(std::array<int64_t, 4> 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<std::unique_ptr<test_case>> 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<<i) {
@@ -11594,6 +11680,24 @@ static std::vector<std::unique_ptr<test_case>> 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 }));