Experiment_1

GGUF

conversational

Model card Files Files and versions

xet

Community

Experiment_1 / openpangu-r-llama.cpp-v0.7.0.patch

HelloSun

Duplicate from HelloSun/openPangu-R-7B-2512-Q4KM-gguf

bb7b6d5 9 days ago

Raw History Blame Contribute Delete

31.9 kB

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, const char *> 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, const char *> 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, llm_tensor_info> 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<llm_graph_context> 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<ggml_tensor *>      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<llm_graph_context> 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<ggml_tensor *> 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<ggml_tensor *> 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<llm_graph_context> llama_model_openpangu_r::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*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<const llama_kv_cache_iswa_context *>(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<const llama_model_openpangu_r &>(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<llm_graph_input_openpangu_r_state>(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;
+}