convert/model/quantize: store Zamba2's per-layer MLP adapter instead of baking it (engine #1036 cause 2)

conversion/zamba2.py

@@ -64,6 +64,15 @@ def __init__(self, *args, **kwargs):
         self._adapter_cache: dict[str, Tensor] = {}
         self._use_attn_adapters = self.hparams.get("use_shared_attention_adapter", False)
         self._use_mlp_adapters = self.hparams.get("use_shared_mlp_adapter", True)
+        # The MLP adapter is stored instead of merged when the checkpoint actually carries one
+        self._store_mlp_adapter = self._use_mlp_adapters and any(
+            "gate_up_proj_adapter_list" in t for t in self.model_tensors)
+        self._adapter_rank = int(self.hparams.get("adapter_rank", 0) or 0)
+        if self._store_mlp_adapter and self._adapter_rank <= 0:
+            raise ValueError(
+                f"{self.dir_model}: the MLP adapters are stored separately, so the checkpoint must declare "
+                "adapter_rank; without it the low-rank tensors cannot be sized"
+            )
 
         adapter_pattern = re.compile(
             r"model\.layers\.(\d+)\.shared_transformer\."
@@ -85,11 +94,16 @@ def __init__(self, *args, **kwargs):
         # Tails of a shared block that this checkpoint stores no per-layer adapter for are
         # byte-identical in every layer using the block, so they are written once, at the
         # block's source layer, and the loader aliases the other layers to it
-        # (llama_model_zamba2).  A tail with an adapter is per layer after merging and stays
-        # per layer -- that one cannot be shared without changing the graph.
+        # (llama_model_zamba2).  Attention adapters, when a checkpoint has them, are still merged
+        # into the per-layer copies -- that is unchanged.  The MLP adapter is NOT merged: the base
+        # gate_up is shared like the rest of the block and each layer keeps its low-rank pair,
+        # applied in the graph as base + B*(A*x) (engine #1036 cause 2).  Baking it per layer cost
+        # 1.13 B params on the 7B, and its 11 duplicated copies of a 102 MB matrix bought nothing.
         adapted = {tail for tail, atype in self._ADAPTER_MAP.items()
-                   if any(atype in t for t in self.model_tensors)}
+                   if atype.startswith("linear_") and any(atype in t for t in self.model_tensors)}
         self._shared_tails = self._SHARED_TAILS - adapted
+        if self._store_mlp_adapter:
+            self._shared_tails = self._shared_tails | {"feed_forward.gate_up_proj.weight"}
 
     def _get_adapter_contribution(self, src_layer: int, adapter_type: str, visit_idx: int) -> Tensor | None:
         """Compute up @ down for a given adapter, returning the additive correction."""
@@ -160,6 +174,11 @@ def set_gguf_parameters(self):
         self.gguf_writer.add_ssm_time_step_rank(n_heads)
         self.gguf_writer.add_ssm_group_count(self.n_group)
 
+        if self._store_mlp_adapter:
+            # the one key that tells the loader which of the two forms this file is: > 0 means ffn_up is the
+            # shared base and every hybrid layer carries a rank-sized A/B pair
+            self.gguf_writer.add_zamba2_ffn_gate_up_adapter_rank(self._adapter_rank)
+
         self.gguf_writer.add_file_type(self.ftype)
 
     # Block tails that carry no per-layer adapter in a Zamba2 checkpoint, i.e. the ones the
@@ -183,7 +202,27 @@ def set_gguf_parameters(self):
     }
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # Skip adapter weights — they are pre-merged into shared weights below
+        # The per-layer MLP adapter is written as its two low-rank factors, at the hybrid layer that owns it,
+        # instead of being baked into that layer's copy of the shared gate_up matrix. The index in the
+        # checkpoint's adapter_list is the hybrid visit index, which is exactly how the layers map to blocks.
+        if self._store_mlp_adapter and "feed_forward.gate_up_proj_adapter_list" in name:
+            m = re.match(
+                r"model\.layers\.(\d+)\.shared_transformer\.feed_forward\.gate_up_proj_adapter_list\.(\d+)\.([01])\.weight$",
+                name)
+            if m is None or bid is None or bid not in self._shared_block_src:
+                # entries for other blocks are nn.Identity and carry no weights; duplicates from ties are
+                # read once, at the block's source layer, like the shared tails above
+                return
+            visit_idx = int(m.group(2))
+            seq = int(m.group(3))  # 0 = down (hidden -> rank), 1 = up (rank -> 2*ff)
+            if visit_idx >= len(self._hybrid_layers):
+                raise ValueError(f"{name}: adapter index {visit_idx} has no hybrid layer")
+            target_bid = self._hybrid_layers[visit_idx]
+            tensor = gguf.MODEL_TENSOR.FFN_GATE_UP_ADAPTER_A if seq == 0 else gguf.MODEL_TENSOR.FFN_GATE_UP_ADAPTER_B
+            yield (self.format_tensor_name(tensor, target_bid), data_torch)
+            return
+
+        # Skip adapter weights that are pre-merged into shared weights below
         if "adapter" in name.lower():
             return
 
@@ -210,14 +249,13 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter
                     yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_POST_NORM, target_bid), data_torch)
                 return
 
-            # Check if this weight has a corresponding adapter type
+            # Check if this weight has a corresponding adapter type. Only the attention adapters are merged
+            # here; the MLP one is stored, so its base is shared and untouched.
             adapter_type = self._ADAPTER_MAP.get(tail)
             use_adapter = False
             if adapter_type:
                 if adapter_type.startswith("linear_") and self._use_attn_adapters:
                     use_adapter = True
-                elif adapter_type == "gate_up_proj_adapter_list" and self._use_mlp_adapters:
-                    use_adapter = True
 
             # Delegate to the parent's tensor-name mapping, once per target layer.
             for visit_local, target_bid in enumerate(target_layers):

gguf-py/gguf/constants.py

@@ -273,6 +273,10 @@ class SSM:
     class KDA:
         HEAD_DIM = "{arch}.kda.head_dim"
 
+    class Zamba2:
+        # names the arch explicitly: 0 or absent = the per-layer MLP adapter is baked into ffn_up
+        FFN_GATE_UP_ADAPTER_RANK = "zamba2.ffn_gate_up_adapter_rank"
+
     class WKV:
         HEAD_SIZE = "{arch}.wkv.head_size"
 
@@ -742,6 +746,8 @@ class MODEL_TENSOR(IntEnum):
     ZAYA_VLORA_DOWN_EXPS_A = auto() # Zaya1-VL
     ZAYA_VLORA_DOWN_EXPS_B = auto() # Zaya1-VL
     SSM_MIX              = auto() # zamba2
+    FFN_GATE_UP_ADAPTER_A = auto() # zamba2: per-layer low-rank MLP adapter, hidden -> rank
+    FFN_GATE_UP_ADAPTER_B = auto() # zamba2: per-layer low-rank MLP adapter, rank -> 2*ff
     TIME_MIX_W0          = auto()
     TIME_MIX_W1          = auto()
     TIME_MIX_W2          = auto()
@@ -1396,6 +1402,8 @@ class MODEL_TENSOR(IntEnum):
     MODEL_TENSOR.SSM_NORM:                  "blk.{bid}.ssm_norm",
     MODEL_TENSOR.SSM_OUT:                   "blk.{bid}.ssm_out",
     MODEL_TENSOR.SSM_MIX:                   "blk.{bid}.ssm_mix",
+    MODEL_TENSOR.FFN_GATE_UP_ADAPTER_A:     "blk.{bid}.ffn_gate_up_adapter_a",
+    MODEL_TENSOR.FFN_GATE_UP_ADAPTER_B:     "blk.{bid}.ffn_gate_up_adapter_b",
     MODEL_TENSOR.SSM_ALPHA:                 "blk.{bid}.ssm_alpha",            # qwen3.5
     MODEL_TENSOR.SSM_BETA_ALPHA:            "blk.{bid}.ssm_ba",
     MODEL_TENSOR.SSM_CONV1D_Q:              "blk.{bid}.ssm_conv1d_q",         # Kimi Linear
@@ -3316,6 +3324,8 @@ class MODEL_TENSOR(IntEnum):
         MODEL_TENSOR.FFN_NORM,
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_UP_ADAPTER_A,
+        MODEL_TENSOR.FFN_GATE_UP_ADAPTER_B,
     ],
     MODEL_ARCH.ZAMBA: [
         MODEL_TENSOR.TOKEN_EMBD,

gguf-py/gguf/gguf_writer.py

@@ -1143,6 +1143,9 @@ def add_ssm_time_step_rank(self, value: int) -> None:
     def add_ssm_group_count(self, value: int) -> None:
         self.add_uint32(Keys.SSM.GROUP_COUNT.format(arch=self.arch), value)
 
+    def add_zamba2_ffn_gate_up_adapter_rank(self, value: int) -> None:
+        self.add_uint32(Keys.Zamba2.FFN_GATE_UP_ADAPTER_RANK, value)
+
     def add_ssm_dt_b_c_rms(self, value: bool) -> None:
         self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value)
 

src/llama-arch.cpp

@@ -412,6 +412,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_XIELU_ALPHA_P,         "xielu.alpha_p"         },
     { LLM_KV_XIELU_BETA,            "xielu.beta"            },
     { LLM_KV_XIELU_EPS,             "xielu.eps"             },
+    { LLM_KV_ZAMBA2_FFN_GATE_UP_ADAPTER_RANK, "zamba2.ffn_gate_up_adapter_rank" },
 
     // deprecated
     { LLM_KV_TOKENIZER_PREFIX_ID, "tokenizer.ggml.prefix_token_id" },
@@ -481,6 +482,8 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
     { LLM_TENSOR_SSM_NORM,                               "blk.%d.ssm_norm" },
     { LLM_TENSOR_SSM_OUT,                                "blk.%d.ssm_out" },
     { LLM_TENSOR_SSM_MIX,                                "blk.%d.ssm_mix" },
+    { LLM_TENSOR_FFN_GATE_UP_ADAPTER_A,                  "blk.%d.ffn_gate_up_adapter_a" },
+    { LLM_TENSOR_FFN_GATE_UP_ADAPTER_B,                  "blk.%d.ffn_gate_up_adapter_b" },
     { LLM_TENSOR_ROPE_FACTORS_LONG,                      "rope_factors_long" },
     { LLM_TENSOR_ROPE_FACTORS_SHORT,                     "rope_factors_short" },
     { LLM_TENSOR_SSM_X,                                  "blk.%d.ssm_x" },
@@ -828,6 +831,8 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_SSM_DT,                     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_SSM_OUT,                    {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_SSM_MIX,                    {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_FFN_GATE_UP_ADAPTER_A,      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_FFN_GATE_UP_ADAPTER_B,      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_SSM_ALPHA,                  {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_SSM_BETA_ALPHA,             {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_TIME_MIX_W1,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},

src/llama-arch.h

@@ -412,6 +412,10 @@ enum llm_kv {
     LLM_KV_XIELU_BETA,
     LLM_KV_XIELU_EPS,
 
+    // zamba2: 0 or absent = the per-layer MLP adapters are baked into ffn_up (files written before this key
+    // existed); > 0 = ffn_up is the shared base and each hybrid layer carries a low-rank adapter pair of this rank
+    LLM_KV_ZAMBA2_FFN_GATE_UP_ADAPTER_RANK,
+
     // deprecated:
     LLM_KV_TOKENIZER_PREFIX_ID,
     LLM_KV_TOKENIZER_SUFFIX_ID,
@@ -513,6 +517,8 @@ enum llm_tensor {
     LLM_TENSOR_SSM_NORM,
     LLM_TENSOR_SSM_OUT,
     LLM_TENSOR_SSM_MIX,  // zamba2: linear mixing between transformer and mamba
+    LLM_TENSOR_FFN_GATE_UP_ADAPTER_A, // zamba2: per-layer low-rank MLP adapter, hidden -> adapter_rank
+    LLM_TENSOR_FFN_GATE_UP_ADAPTER_B, // zamba2: per-layer low-rank MLP adapter, adapter_rank -> 2*n_ff
     LLM_TENSOR_SSM_BETA_ALPHA,      // qwen3next
     LLM_TENSOR_SSM_ALPHA,           // qwen3.5
     // Kimi Linear KDA (using SSM_ prefix for consistency)

src/llama-graph.cpp

@@ -1602,7 +1602,8 @@ ggml_tensor * llm_graph_context::build_ffn(
          ggml_tensor * act_scales,
      llm_ffn_op_type   type_op,
    llm_ffn_gate_type   type_gate,
-                 int   il) const {
+                 int   il,
+         ggml_tensor * up_delta) const {
     // NVFP4 support is currently restricted to
     // 1) LORA absence (*_s would be applied after LORA residual, which is incorrect)
     // 2) bias absense (*_s would be applied after bias addition, which is incorrect)
@@ -1634,6 +1635,13 @@ ggml_tensor * llm_graph_context::build_ffn(
         cb(tmp, "ffn_up_b", il);
     }
 
+    // a low-rank adapter added to the fused gate/up output (zamba2: base + B*(A*x)); an addition to the matmul
+    // result, so it is placed before the split and the activation, exactly where the checkpoint applies it
+    if (up_delta) {
+        tmp = ggml_add(ctx0, tmp, up_delta);
+        cb(tmp, "ffn_up_delta", il);
+    }
+
     if (up_s) {
         tmp = ggml_mul(ctx0, tmp, up_s);
         cb(tmp, "ffn_up_s", il);

src/llama-graph.h

@@ -1005,7 +1005,8 @@ struct llm_graph_context {
              ggml_tensor * act_scales,
          llm_ffn_op_type   type_op,
        llm_ffn_gate_type   type_gate,
-                     int   il) const;
+                     int   il,
+             ggml_tensor * up_delta = nullptr) const;
 
     // build MoE FFN without bias tensors
     ggml_tensor * build_moe_ffn(

src/llama-hparams.h

@@ -167,6 +167,10 @@ struct llama_hparams {
     uint32_t ssm_dt_rank = 0;
     uint32_t ssm_n_group = 0;
 
+    // zamba2: 0 or absent = the per-layer MLP adapter is baked into ffn_up; > 0 = ffn_up is the shared base and
+    // each hybrid layer carries a low-rank adapter pair of this rank (engine #1036 cause 2)
+    uint32_t zamba2_ffn_gate_up_adapter_rank = 0;
+
     // for Kimi Linear KDA
     uint32_t n_embd_head_kda = 0;
 

src/llama-model-saver.cpp

@@ -401,6 +401,11 @@ void llama_model_saver::add_kv_from_model() {
     } else {
         add_kv(LLM_KV_SSM_GROUP_COUNT,               hparams.ssm_n_group);
     }
+    if (hparams.zamba2_ffn_gate_up_adapter_rank > 0) {
+        // only written when the tensors it describes are present, so the saver cannot claim the adapter form
+        // for a file that has the adapter baked into ffn_up
+        add_kv(LLM_KV_ZAMBA2_FFN_GATE_UP_ADAPTER_RANK, hparams.zamba2_ffn_gate_up_adapter_rank);
+    }
     add_kv(LLM_KV_SSM_DT_B_C_RMS,                    hparams.ssm_dt_b_c_rms);
 
     add_kv(LLM_KV_KDA_HEAD_DIM,                      hparams.n_embd_head_kda);

src/llama-model.h

@@ -303,6 +303,10 @@ struct llama_layer {
     struct ggml_tensor * ffn_gate_enc = nullptr;
     struct ggml_tensor * ffn_down_enc = nullptr;
     struct ggml_tensor * ffn_up_enc   = nullptr;
+    // zamba2: the per-layer MLP adapter (hidden -> rank -> 2*n_ff), applied as base + B*(A*x); null when the
+    // adapter was baked into ffn_up at conversion time
+    struct ggml_tensor * ffn_gate_up_adapter_a = nullptr;
+    struct ggml_tensor * ffn_gate_up_adapter_b = nullptr;
 
     // ff MoE
     struct ggml_tensor * ffn_gate_inp      = nullptr;

src/llama-quant.cpp

@@ -680,6 +680,12 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type
     return new_type;
 }
 
+// zamba2's per-layer MLP adapter pair (engine #1036 cause 2). Matched by name because the pair is an arch-level
+// tensor a quantize rule needs to see before the type is chosen, and it has no category of its own.
+static bool is_zamba2_mlp_adapter(const char * name) {
+    return std::strstr(name, "ffn_gate_up_adapter_") != nullptr;
+}
+
 // outer wrapper: determine the ggml_type that this tensor should be quantized to
 static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_model_quantize_params * params, const ggml_tensor * tensor, ggml_type default_type, const tensor_metadata & tm) {
     if (!tensor_allows_quantization(params, qs.model.arch, tensor)) {
@@ -728,8 +734,19 @@ static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_mod
         }
 
         // if not manual - use the standard logic for choosing the quantization type based on the selected mixture
-        if (!manual && !params->pure) {
-            new_type = llama_tensor_get_type_impl(qs, new_type, tensor, params->ftype, tm.category);
+        if (!manual) {
+            if (!params->pure && is_zamba2_mlp_adapter(tensor->name)) {
+                // zamba2's per-layer MLP adapter is a rank-128 pair applied at run time (base + B*(A*x)). Quantizing
+                // it saves little and costs the whole pair its place on the device: the B contraction has
+                // input_size = rank = 128, which is below the block formats' `input_size >= 256` minimum (their
+                // own block layout) while the float path accepts `>= 4 && % 4 == 0`. Measured on a 2.7B Q5_K_M,
+                // quantizing the pair puts 9 of 3051 nodes on the CPU (one per hybrid layer) and takes prefill from
+                // 589.5 to 463.7 tok/s; keeping it at F16 costs 31.8 MiB of 1804.9. --pure still forces the
+                // requested type on every tensor, and an explicit --tensor-type wins over this default.
+                new_type = GGML_TYPE_F16;
+            } else if (!params->pure) {
+                new_type = llama_tensor_get_type_impl(qs, new_type, tensor, params->ftype, tm.category);
+            }
         }
 
         // incompatible tensor shapes are handled here - fallback to a compatible type