diff --git a/conversion/__init__.py b/conversion/__init__.py index 8ce4cb7..f7455f4 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", + "PanguEmbeddedRForCausalLM": "pangu", "Phi3ForCausalLM": "phi", "Phi4ForCausalLMV": "phi", "PhiForCausalLM": "phi", diff --git a/conversion/base.py b/conversion/base.py index a5a80f2..1a5276b 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -3121,6 +3121,12 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st arch = text_config["architectures"][0] elif model_type == ModelType.MMPROJ and vision_config.get("architectures") is not None: arch = vision_config["architectures"][0] + + # openPangu-R reuses the PanguEmbedded architecture name, but its config carries the + # sliding-window / attention-sink / gate / conv extensions: route it to its own class + if model_type == ModelType.TEXT and arch == "PanguEmbeddedForCausalLM" and "param_sink_number" in hparams: + arch = "PanguEmbeddedRForCausalLM" + if arch is None: raise ValueError("Failed to detect model architecture") return arch diff --git a/conversion/pangu.py b/conversion/pangu.py index 74c7653..14b2250 100644 --- a/conversion/pangu.py +++ b/conversion/pangu.py @@ -45,3 +45,71 @@ class PanguEmbeddedModel(TextModel): logger.info("Skipping tied output layer 'lm_head.weight'") return yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("PanguEmbeddedRForCausalLM") +@ModelBase.example("openpangu/openPangu-R-7B-2512") +class PanguEmbeddedRModel(TextModel): + # openPangu-R: PanguEmbedded config + interleaved sliding-window/full attention, + # per-layer attention sinks (learned K/V prefix), gated attention with per-head RMSNorm, + # partial RoPE, a causal depthwise conv merge in the first and last layers, and an MTP head + # (layer index == block_count, not used for inference). + model_arch = gguf.MODEL_ARCH.PANGU_EMBED_R + + def set_vocab(self): + self._set_vocab_sentencepiece() + + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + if tokenizer_config_file.is_file(): + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + tokenizer_config_json = json.load(f) + if "add_prefix_space" in tokenizer_config_json: + self.gguf_writer.add_add_space_prefix(tokenizer_config_json["add_prefix_space"]) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + + # partial RoPE: only rotary_percent of head_dim is rotated (rounded up to an even count, as in HF) + head_dim = hparams["head_dim"] + rot_dim = int(head_dim * hparams["rotary_percent"]) + rot_dim += rot_dim % 2 + self.gguf_writer.add_rope_dimension_count(rot_dim) + + self.gguf_writer.add_sliding_window(hparams["sliding_window"]) + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hparams["layer_types"]]) + + def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: + # sinks and the conv merge are small; keep them in F32 + if any(self.match_model_tensor_name(new_name, key, bid) for key in ( + gguf.MODEL_TENSOR.ATTN_SINK_K, + gguf.MODEL_TENSOR.ATTN_SINK_V, + gguf.MODEL_TENSOR.FFN_MERGE_CONV, + )): + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # the MTP module (layer block_count) and the second final norm are unused for inference + if (bid is not None and bid >= self.block_count) or name == "model.norms.1.weight": + return + + # HF param_sink_key/value: [n_sink, n_kv_head, head_dim] + # GGUF ne = [head_dim, n_sink, n_kv_head] (key) and [n_sink, head_dim, n_kv_head] (value) + if name.endswith(".self_attn.param_sink_key"): + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_SINK_K, bid), + data_torch.permute(1, 0, 2).contiguous()) + return + if name.endswith(".self_attn.param_sink_value"): + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_SINK_V, bid), + data_torch.permute(1, 2, 0).contiguous()) + return + + # HF depthwise Conv1d weight [n_embd, 1, 3] -> GGUF ne = [n_embd, 3] + if name.endswith(".merge_conv.weight"): + yield (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_MERGE_CONV, bid), + data_torch.squeeze(1).permute(1, 0).contiguous()) + return + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a39d046..8a7f15a 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -674,6 +674,7 @@ class MODEL_ARCH(IntEnum): MINIMAXM3 = auto() RND1 = auto() PANGU_EMBED = auto() + PANGU_EMBED_R = auto() MISTRAL3 = auto() EAGLE3 = auto() DFLASH = auto() @@ -743,6 +744,9 @@ class MODEL_TENSOR(IntEnum): ATTN_ROT_EMBD = auto() ATTN_SINKS = auto() ATTN_GATE = auto() + ATTN_SINK_K = auto() + ATTN_SINK_V = auto() + FFN_MERGE_CONV = auto() FFN_GATE_INP = auto() FFN_GATE_INP_SHEXP = auto() FFN_NORM = auto() @@ -1479,6 +1483,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.COGVLM: "cogvlm", MODEL_ARCH.RND1: "rnd1", MODEL_ARCH.PANGU_EMBED: "pangu-embedded", + MODEL_ARCH.PANGU_EMBED_R: "openpangu-r", MODEL_ARCH.MISTRAL3: "mistral3", MODEL_ARCH.EAGLE3: "eagle3", MODEL_ARCH.DFLASH: "dflash", @@ -1544,6 +1549,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd", MODEL_TENSOR.ATTN_SINKS: "blk.{bid}.attn_sinks", MODEL_TENSOR.ATTN_GATE: "blk.{bid}.attn_gate", + MODEL_TENSOR.ATTN_SINK_K: "blk.{bid}.attn_sink_k", + MODEL_TENSOR.ATTN_SINK_V: "blk.{bid}.attn_sink_v", + MODEL_TENSOR.FFN_MERGE_CONV: "blk.{bid}.ffn_merge_conv", MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm", MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm", MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm", @@ -5488,6 +5496,25 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.PANGU_EMBED_R: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.ATTN_SINK_K, + MODEL_TENSOR.ATTN_SINK_V, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_MERGE_CONV, + ], MODEL_ARCH.MISTRAL3: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 63301cc..df154b8 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -101,6 +101,7 @@ class TensorNameMap: "gpt_neox.final_layer_norm", # gptneox "transformer.ln_f", # gpt2 gpt-j falcon jais exaone "model.norm", # llama-hf baichuan internlm2 olmoe olmo2 phimoe plamo2 + "model.norms.0", # openpangu-r (final norm) "norm", # llama-pth "transformer.norm_f", # mpt dbrx "ln_f", # refact bloom qwen gpt2 @@ -409,6 +410,7 @@ class TensorNameMap: "encoder.layers.{bid}.norm1", # nomic-bert "transformer.decoder_layer.{bid}.rms_norm_1", # Grok "model.layers.{bid}.post_attn_norm", # grok-2 + "model.layers.{bid}.self_attn.groupnorm", # openpangu-r (per-head RMSNorm before gate) "transformer.blocks.{bid}.norm_attn_norm.norm_2", # dbrx ), @@ -435,6 +437,7 @@ class TensorNameMap: MODEL_TENSOR.ATTN_GATE: ( "model.layers.{bid}.self_attn.gate_proj", # afmoe muse-glimmer + "model.layers.{bid}.self_attn.attention_gate", # openpangu-r "model.layers.{bid}.linear_attn.in_proj_z", # qwen3.5 "layers.{bid}.linear_attn.in_proj_z", # qwen3.5 text "model.layers.{bid}.self_attn.g_proj", # step3.5 head-wise attention gate diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index d1dd4cb..d30be8c 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -144,6 +144,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_COGVLM, "cogvlm" }, { LLM_ARCH_RND1, "rnd1" }, { LLM_ARCH_PANGU_EMBED, "pangu-embedded" }, + { LLM_ARCH_OPENPANGU_R, "openpangu-r" }, { LLM_ARCH_MISTRAL3, "mistral3" }, { LLM_ARCH_EAGLE3, "eagle3" }, { LLM_ARCH_DFLASH, "dflash" }, @@ -500,6 +501,9 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, { LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" }, + { LLM_TENSOR_ATTN_SINK_K, "blk.%d.attn_sink_k" }, + { LLM_TENSOR_ATTN_SINK_V, "blk.%d.attn_sink_v" }, + { LLM_TENSOR_FFN_MERGE_CONV, "blk.%d.ffn_merge_conv" }, { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, { LLM_TENSOR_FFN_POST_NORM_1, "blk.%d.post_ffw_norm_1" }, { LLM_TENSOR_FFN_POST_NORM_2, "blk.%d.post_ffw_norm_2" }, @@ -839,6 +843,9 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_ATTN_QKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_OUT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_SINK_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_SINK_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_MERGE_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_FFN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_FFN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 4e8e534..0fa9b42 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -146,6 +146,7 @@ enum llm_arch { LLM_ARCH_COGVLM, LLM_ARCH_RND1, LLM_ARCH_PANGU_EMBED, + LLM_ARCH_OPENPANGU_R, LLM_ARCH_MISTRAL3, LLM_ARCH_MISTRAL4, LLM_ARCH_PADDLEOCR, @@ -500,6 +501,9 @@ enum llm_tensor { LLM_TENSOR_ATTN_ROT_EMBD, LLM_TENSOR_ATTN_SINKS, LLM_TENSOR_ATTN_GATE, + LLM_TENSOR_ATTN_SINK_K, + LLM_TENSOR_ATTN_SINK_V, + LLM_TENSOR_FFN_MERGE_CONV, LLM_TENSOR_FFN_GATE_INP, LLM_TENSOR_FFN_GATE_INP_SHEXP, LLM_TENSOR_FFN_NORM, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index ded790b..a811c9d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -324,6 +324,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_cogvlm(params); case LLM_ARCH_PANGU_EMBED: return new llama_model_pangu_embed(params); + case LLM_ARCH_OPENPANGU_R: + return new llama_model_openpangu_r(params); case LLM_ARCH_QWEN3NEXT: return new llama_model_qwen3next(params); case LLM_ARCH_QWEN35: @@ -3590,6 +3592,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_COGVLM: case LLM_ARCH_PANGU_EMBED: + case LLM_ARCH_OPENPANGU_R: case LLM_ARCH_AFMOE: case LLM_ARCH_LAGUNA: case LLM_ARCH_QWEN3NEXT: diff --git a/src/llama-model.h b/src/llama-model.h index bcfbf1e..4ac2432 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -529,6 +529,12 @@ struct llama_layer { // openai-moe struct ggml_tensor * attn_sinks = nullptr; + // openpangu-r + struct ggml_tensor * attn_gate = nullptr; + struct ggml_tensor * attn_sink_k = nullptr; + struct ggml_tensor * attn_sink_v = nullptr; + struct ggml_tensor * ffn_merge_conv = nullptr; + // DeepSeek-V4 struct ggml_tensor * attn_kv_norm = nullptr; struct ggml_tensor * hc_attn_fn = nullptr; diff --git a/src/models/models.h b/src/models/models.h index 88fd07d..773e3fe 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2378,6 +2378,42 @@ struct llama_model_pangu_embed : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +// openPangu-R-7B-2512 +struct llama_model_openpangu_r : public llama_model_base { + llama_model_openpangu_r(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; + + // persistent history of the merge_conv layers: the last 2 post-attention hidden states [n_embd, 2] + // (indexed by layer, nullptr elsewhere); owned by the model, zeroed when a sequence restarts at pos 0 + ggml_context_ptr conv_state_ctx; + ggml_backend_buffer_ptr conv_state_buf; + std::vector conv_state; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + // causal attention with per-layer learned sink keys/values (not cached, no RoPE) + ggml_tensor * build_attn_sinks( + llm_graph_input_attn_kv_iswa * inp, + const llama_layer & layer, + ggml_tensor * attn_in, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + int il); + + // depthwise causal conv over tokens (kernel 3), continuing from the persisted history + ggml_tensor * build_merge_conv( + ggml_tensor * h, + const llama_layer & layer, + ggml_tensor * state, + int il); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_qwen3next : public llama_model_base { llama_model_qwen3next(const struct llama_model_params & params) : llama_model_base(params) {} diff --git a/src/models/openpangu-r.cpp b/src/models/openpangu-r.cpp new file mode 100644 index 0000000..5ece53e --- /dev/null +++ b/src/models/openpangu-r.cpp @@ -0,0 +1,366 @@ +#include "models.h" +#include "llama-kv-cache-iswa.h" + +// openPangu-R-7B-2512 (HF architecture "PanguEmbeddedForCausalLM" with the R extensions) +// +// per layer: +// h1 = h + wo( groupnorm(attn(q, k, v)) * sigmoid(attn_gate(x)) ) x = rmsnorm(h) +// x' = merge_conv(h1) (layers 0 and N-1 only; causal depthwise conv, kernel 3, no bias) +// h2 = h1 + ffn(rmsnorm(x')) +// attention: +// - interleaved sliding-window (even layers, window 128) and full attention (odd layers) +// - partial RoPE (first 32 of 128 dims, NEOX layout) +// - per-layer learned sinks: n_sink (key, value) pairs per KV head that every query can attend to; +// they are not cached and not rotated, so they are added to the softmax as extra logits +// - per-head RMSNorm (groupnorm, weight shared across heads) on the attention output +// the MTP module (layer N) is not used for inference +// +// limitations: one sequence per batch; the merge_conv history is kept in the model, so it only follows +// a single sequence that is decoded contiguously (positions restart at 0 for a new sequence) + +namespace { + +// zeroes the merge_conv history when a sequence starts at position 0 +class llm_graph_input_openpangu_r_state : public llm_graph_input_i { +public: + explicit llm_graph_input_openpangu_r_state(std::vector states) : states(std::move(states)) {} + + void set_input(const llama_ubatch * ubatch) override { + if (ubatch->n_seqs_unq > 1) { + GGML_ABORT("openpangu-r: only one sequence per batch is supported"); + } + if (ubatch->pos != nullptr && ubatch->n_tokens > 0 && ubatch->pos[0] == 0) { + for (ggml_tensor * s : states) { + if (s != nullptr) { + ggml_backend_tensor_memset(s, 0, 0, ggml_nbytes(s)); + } + } + } + } + + std::vector states; +}; + +} // namespace + +void llama_model_openpangu_r::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + // interleaved sliding-window / full attention; the pattern comes from the GGUF (sliding_window_pattern) + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + load_swa_pattern(ml, 2); + + // the HF implementation uses a single rotary embedding for both layer types + ml.get_key(LLM_KV_ROPE_FREQ_BASE, hparams.rope_freq_base_train_swa); + + switch (hparams.n_layer()) { + case 27: type = LLM_TYPE_7B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_openpangu_r::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + // output + 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}, TENSOR_NOT_REQUIRED); + + // if output is NULL, init from the input tok embed + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + // number of sink tokens: dim 1 of attn_sink_k (GGUF ne = [head_dim, n_sink, n_kv_head]) + const ggml_tensor * sink_meta = ml.get_tensor_meta(tn(LLM_TENSOR_ATTN_SINK_K, "weight", 0).str().c_str()); + if (sink_meta == nullptr) { + throw std::runtime_error("openpangu-r: missing attn_sink_k tensor"); + } + const int64_t n_sink = sink_meta->ne[1]; + + conv_state.assign(n_layer, nullptr); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, 0); + + layer.attn_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); + layer.attn_out_norm = create_tensor(tn(LLM_TENSOR_ATTN_OUT_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_sink_k = create_tensor(tn(LLM_TENSOR_ATTN_SINK_K, "weight", i), {n_embd_head_k, n_sink, n_head_kv}, 0); + layer.attn_sink_v = create_tensor(tn(LLM_TENSOR_ATTN_SINK_V, "weight", i), {n_sink, n_embd_head_v, n_head_kv}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + + if (i == 0 || i == n_layer - 1) { + layer.ffn_merge_conv = create_tensor(tn(LLM_TENSOR_FFN_MERGE_CONV, "weight", i), {n_embd, 3}, 0); + } + } + + // persistent merge_conv history (2 tokens x n_embd per conv layer), allocated on the CPU + int n_conv = 0; + for (int i = 0; i < n_layer; ++i) { + n_conv += layers[i].ffn_merge_conv != nullptr; + } + + ggml_init_params params = { + /*.mem_size =*/ ggml_tensor_overhead() * (n_conv + 1), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + conv_state_ctx.reset(ggml_init(params)); + if (!conv_state_ctx) { + throw std::runtime_error("openpangu-r: failed to create the conv state context"); + } + + for (int i = 0; i < n_layer; ++i) { + if (layers[i].ffn_merge_conv == nullptr) { + continue; + } + conv_state[i] = ggml_new_tensor_2d(conv_state_ctx.get(), GGML_TYPE_F32, n_embd, 2); + ggml_format_name(conv_state[i], "openpangu_r_conv_state_%d", i); + } + + conv_state_buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(conv_state_ctx.get(), ggml_backend_cpu_buffer_type())); + if (!conv_state_buf) { + throw std::runtime_error("openpangu-r: failed to allocate the conv state buffer"); + } +} + +std::unique_ptr llama_model_openpangu_r::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// the attention of one layer; sinks are extra softmax terms, the K/V cache is written as usual +ggml_tensor * llama_model_openpangu_r::graph::build_attn_sinks( + llm_graph_input_attn_kv_iswa * inp, + const llama_layer & layer, + ggml_tensor * attn_in, + ggml_tensor * q, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + int il) { + const bool is_swa = hparams.is_swa(il); + + GGML_ASSERT(!(is_swa ? inp->self_k_rot_swa : inp->self_k_rot) && "openpangu-r: Hadamard-rotated KV cache is not supported"); + GGML_ASSERT(!(is_swa ? inp->self_v_rot_swa : inp->self_v_rot) && "openpangu-r: Hadamard-rotated KV cache is not supported"); + + const int64_t n_embd_head = q->ne[0]; + const int64_t n_tokens = q->ne[2]; + const int64_t n_sink = layer.attn_sink_k->ne[1]; + const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); + + const auto * mctx_iswa = static_cast(mctx); + const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base(); + + // store the current K/V, then read the whole cache (same as llm_graph_context::build_attn) + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, is_swa ? inp->get_k_idxs_swa() : inp->get_k_idxs(), il)); + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, is_swa ? inp->get_v_idxs_swa() : inp->get_v_idxs(), il)); + + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = mctx_cur->get_v(ctx0, il); + + GGML_ASSERT(k->ne[3] == 1 && "openpangu-r: multi-stream KV cache is not supported"); + + const bool v_trans = v->nb[1] > v->nb[2]; + + // [D, n_tokens, n_head], [D, n_kv, n_kv_head], [n_kv, D, n_kv_head] + ggml_tensor * qp = ggml_permute(ctx0, q, 0, 2, 1, 3); + ggml_tensor * kp = ggml_permute(ctx0, k, 0, 2, 1, 3); + ggml_tensor * vp = ggml_permute(ctx0, v, 0, 2, 1, 3); + if (!v_trans) { + vp = ggml_cont(ctx0, ggml_transpose(ctx0, vp)); + } + + // cached keys: [n_kv, n_tokens, n_head] + ggml_tensor * kq = ggml_mul_mat(ctx0, kp, qp); + kq = ggml_scale(ctx0, kq, kq_scale); + + ggml_tensor * kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask(); + if (kq_mask != nullptr) { + ggml_tensor * mask = ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); + mask = ggml_cont(ctx0, ggml_view_3d(ctx0, mask, kq->ne[0], n_tokens, 1, mask->nb[1], mask->nb[2], 0)); + kq = ggml_add(ctx0, kq, mask); + } + + // sink keys: not cached, not rotated, always visible: [n_sink, n_tokens, n_head] + ggml_tensor * ks = ggml_mul_mat(ctx0, layer.attn_sink_k, qp); + ks = ggml_scale(ctx0, ks, kq_scale); + + // one softmax over [cached keys | sinks] + ggml_tensor * p = ggml_soft_max(ctx0, ggml_concat(ctx0, kq, ks, 0)); + + const int64_t n_kv = kq->ne[0]; + // the views have row strides of the full [n_kv + n_sink] row; the CPU mul_mat needs contiguous rows + ggml_tensor * p_kv = ggml_cont(ctx0, ggml_view_3d(ctx0, p, n_kv, n_tokens, kq->ne[2], p->nb[1], p->nb[2], 0)); + ggml_tensor * p_sk = ggml_cont(ctx0, ggml_view_3d(ctx0, p, n_sink, n_tokens, kq->ne[2], p->nb[1], p->nb[2], n_kv*p->nb[0])); + + ggml_tensor * o_kv = ggml_mul_mat(ctx0, vp, p_kv); // [D, n_tokens, n_head] + ggml_tensor * o_sk = ggml_mul_mat(ctx0, layer.attn_sink_v, p_sk); // [D, n_tokens, n_head] + ggml_tensor * o = ggml_add(ctx0, o_kv, o_sk); + + o = ggml_cont(ctx0, ggml_permute(ctx0, o, 0, 2, 1, 3)); // [D, n_head, n_tokens] + cb(o, "kqv", il); + + // per-head RMSNorm, then the sigmoid gate computed from the attention input + o = build_norm(o, layer.attn_out_norm, NULL, LLM_NORM_RMS, il); + o = ggml_reshape_2d(ctx0, o, n_embd_head * hparams.n_head(il), n_tokens); + + ggml_tensor * gate = build_lora_mm(layer.attn_gate, attn_in, nullptr); + o = ggml_mul(ctx0, o, ggml_sigmoid(ctx0, gate)); + + o = build_lora_mm(layer.wo, o, layer.wo_s); + if (layer.wo_b) { + o = ggml_add(ctx0, o, layer.wo_b); + } + + return o; +} + +// causal depthwise conv over the token axis: out[t] = w0*h[t-2] + w1*h[t-1] + w2*h[t] +// history of the previous two tokens comes from `state`, which is updated with the last two tokens of h +ggml_tensor * llama_model_openpangu_r::graph::build_merge_conv( + ggml_tensor * h, + const llama_layer & layer, + ggml_tensor * state, + int il) { + const int64_t n_embd = h->ne[0]; + const int64_t n_tokens = h->ne[1]; + + GGML_ASSERT(state->ne[0] == n_embd && state->ne[1] == 2); + + // [history(2) | h(n_tokens)] along the token axis + ggml_tensor * cat = ggml_concat(ctx0, state, h, 1); + cb(cat, "merge_conv_in", il); + + // persist the last two tokens for the next batch + ggml_tensor * next = ggml_view_2d(ctx0, cat, n_embd, 2, cat->nb[1], n_tokens*cat->nb[1]); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, next, state)); + + ggml_tensor * W = ggml_cast(ctx0, layer.ffn_merge_conv, GGML_TYPE_F32); // [n_embd, 3] + + ggml_tensor * out = nullptr; + for (int k = 0; k < 3; ++k) { + ggml_tensor * x_k = ggml_view_2d(ctx0, cat, n_embd, n_tokens, cat->nb[1], k*cat->nb[1]); + ggml_tensor * w_k = ggml_view_2d(ctx0, W, n_embd, 1, W->nb[1], k*W->nb[1]); + ggml_tensor * t = ggml_mul(ctx0, x_k, w_k); + out = out == nullptr ? t : ggml_add(ctx0, out, t); + } + + cb(out, "merge_conv_out", il); + return out; +} + +llama_model_openpangu_r::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const auto & self = static_cast(model); + + const int64_t n_embd_head = hparams.n_embd_head_k(); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + // inp_pos - contains the positions + ggml_tensor * inp_pos = build_inp_pos(); + + res->add_input(std::make_unique(self.conv_state)); + + auto * inp_attn = build_attn_inp_kv_iswa(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * inpSA = inpL; + + // norm + cur = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // self attention + { + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv, il); + + // partial RoPE: only the first n_rot dims are rotated + const float freq_base_l = model.get_rope_freq_base (cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = build_attn_sinks(inp_attn, layer, cur, Qcur, Kcur, Vcur, il); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + // merge conv runs over all tokens of the batch, before the output rows are selected + ggml_tensor * x = ffn_inp; + if (layer.ffn_merge_conv != nullptr) { + x = build_merge_conv(ffn_inp, layer, self.conv_state[il], il); + } + + if (il == n_layer - 1 && inp_out_ids) { + x = ggml_get_rows(ctx0, x, inp_out_ids); + ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids); + } + + // feed-forward network + cur = build_norm(x, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + layer.ffn_up, NULL, NULL, + layer.ffn_gate, NULL, NULL, + layer.ffn_down, NULL, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, + model.output_norm, NULL, + LLM_NORM_RMS, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + // lm_head + cur = build_lora_mm(model.output, cur); + + cb(cur, "result_output", -1); + res->t_logits = cur; +}