llama-cpp-minimax-adaptive-gguf-export.patch

diff --git a/conversion/minimax.py b/conversion/minimax.py
index 4857775cb..6e0b78303 100644
--- a/conversion/minimax.py
+++ b/conversion/minimax.py
@@ -7,25 +7,149 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf
+from .base import LazyTorchTensor, ModelBase, TextModel, gguf
 
 
 @ModelBase.register("MiniMaxM2ForCausalLM")
 class MiniMaxM2Model(TextModel):
     model_arch = gguf.MODEL_ARCH.MINIMAXM2
-    _experts_cache: dict[int, dict[str, Tensor]] = {}
+    _experts_cache: dict[int, dict[str, Tensor]]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._experts_cache = {}
+        self._adaptive_retained_by_layer = self._load_adaptive_retained_by_layer()
+        self._uniform_expert_count = self.find_hparam(["num_local_experts", "num_experts"])
+
+    def _load_adaptive_retained_by_layer(self) -> dict[int, int] | None:
+        raw_counts = self.hparams.get("reap_adaptive_retained_experts_by_layer")
+        if not raw_counts:
+            if self.hparams.get("reap_adaptive_nonuniform_experts"):
+                raise ValueError(
+                    "MiniMax-M2 adaptive checkpoint declares reap_adaptive_nonuniform_experts "
+                    "but lacks reap_adaptive_retained_experts_by_layer"
+                )
+            return None
+
+        if not isinstance(raw_counts, dict):
+            raise ValueError(
+                "reap_adaptive_retained_experts_by_layer must be a layer->count object, "
+                f"got {type(raw_counts).__name__}"
+            )
+
+        max_experts = self.find_hparam(["num_local_experts", "num_experts"])
+        counts: dict[int, int] = {}
+        for raw_layer, raw_count in raw_counts.items():
+            layer = int(raw_layer)
+            count = int(raw_count)
+            if layer < 0 or layer >= self.block_count:
+                raise ValueError(
+                    f"adaptive retained expert metadata has out-of-range layer {layer}; "
+                    f"block_count={self.block_count}"
+                )
+            if count < 1 or count > max_experts:
+                raise ValueError(
+                    f"adaptive retained expert count for layer {layer} is {count}; "
+                    f"expected 1..{max_experts}"
+                )
+            counts[layer] = count
+
+        missing = sorted(set(range(self.block_count)) - set(counts))
+        if missing:
+            raise ValueError(
+                "MiniMax-M2 adaptive GGUF export requires counts for every layer; "
+                f"missing layers={missing[:8]}{'...' if len(missing) > 8 else ''}"
+            )
+        if max(counts.values()) != max_experts:
+            raise ValueError(
+                "MiniMax-M2 adaptive config num_local_experts must equal the max retained "
+                f"layer count for padded GGUF export; got config={max_experts} "
+                f"metadata_max={max(counts.values())}"
+            )
+        return counts
+
+    def _layer_expert_count(self, bid: int) -> int:
+        if self._adaptive_retained_by_layer is None:
+            return self._uniform_expert_count
+        return self._adaptive_retained_by_layer[bid]
+
+    def _pad_expert_axis(self, data_torch: Tensor, count: int, *, dim: int, fill_value: float = 0.0) -> Tensor:
+        if count == self._uniform_expert_count:
+            return data_torch
+        if data_torch.shape[dim] != count:
+            raise ValueError(
+                "MiniMax-M2 adaptive tensor shape does not match metadata: "
+                f"shape={tuple(data_torch.shape)} expert_dim={dim} "
+                f"shape_count={data_torch.shape[dim]} metadata_count={count}"
+            )
+
+        pad_shape = list(data_torch.shape)
+        pad_shape[dim] = self._uniform_expert_count - count
+        if isinstance(data_torch, LazyTorchTensor):
+            pad = LazyTorchTensor(
+                meta=LazyTorchTensor.meta_with_dtype_and_shape(data_torch.dtype, tuple(pad_shape)),
+                args=(data_torch,),
+                func=lambda source: torch.full(
+                    tuple(pad_shape),
+                    fill_value,
+                    dtype=source.dtype,
+                    device=source.device,
+                ),
+            )
+        else:
+            pad = torch.full(
+                tuple(pad_shape),
+                fill_value,
+                dtype=data_torch.dtype,
+                device=data_torch.device,
+            )
+        return torch.cat([data_torch, pad], dim=dim).contiguous()
+
+    def _add_adaptive_padding_metadata(self):
+        if self._adaptive_retained_by_layer is None:
+            return
+        counts = [self._adaptive_retained_by_layer[layer] for layer in range(self.block_count)]
+        self.gguf_writer.add_bool("minimax-m2.reap_adaptive_padded_uniform", True)
+        self.gguf_writer.add_bool("minimax-m2.reap_adaptive_dummy_experts_masked", True)
+        self.gguf_writer.add_array("minimax-m2.reap_adaptive_retained_experts_by_layer", counts)
 
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
 
         self.gguf_writer.add_expert_feed_forward_length(self.find_hparam(["intermediate_size"]))
         self.gguf_writer.add_rope_dimension_count(self.find_hparam(["rotary_dim"]))
+        self._add_adaptive_padding_metadata()
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        if self._experts_cache:
+            pending = {bid: len(cache) for bid, cache in sorted(self._experts_cache.items())}
+            raise RuntimeError(
+                "MiniMax-M2 expert export ended with incomplete per-layer expert caches; "
+                f"pending_tensors_by_layer={pending}"
+            )
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
+        if bid is not None and self._adaptive_retained_by_layer is not None:
+            count = self._layer_expert_count(bid)
+            if name.endswith(".block_sparse_moe.gate.weight"):
+                data_torch = self._pad_expert_axis(data_torch, count, dim=0)
+            elif name.endswith((".block_sparse_moe.e_score_correction.bias", ".block_sparse_moe.e_score_correction_bias")):
+                # Large finite negative instead of -inf: selection scores are
+                # sigmoid(logits) + bias with sigmoid in [0, 1], so -1e9 makes a
+                # padded expert unselectable exactly like -inf would, while
+                # ggml_validate_row_data (llama-quantize) rejects inf values.
+                data_torch = self._pad_expert_axis(
+                    data_torch,
+                    count,
+                    dim=0,
+                    fill_value=-1.0e9,
+                )
+
         # merge expert weights
         if 'experts' in name:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
             assert bid is not None
+            n_experts = self._layer_expert_count(bid)
 
             expert_cache = self._experts_cache.setdefault(bid, {})
             expert_cache[name] = data_torch
@@ -44,6 +168,7 @@ class MiniMaxM2Model(TextModel):
                     del expert_cache[ename]
 
                 data_torch = torch.stack(datas, dim=0)
+                data_torch = self._pad_expert_axis(data_torch, n_experts, dim=0)
                 merged_name = f"model.layers.{bid}.block_sparse_moe.experts.{w_name}.weight"
                 new_name = self.map_tensor_name(merged_name)
                 yield from super().modify_tensors(data_torch, new_name, bid)
diff --git a/tests/test_minimax_adaptive_export.py b/tests/test_minimax_adaptive_export.py
new file mode 100644
index 000000000..0a694bafc
--- /dev/null
+++ b/tests/test_minimax_adaptive_export.py
@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from conversion.base import LazyTorchTensor, TextModel, gguf
+from conversion.minimax import MiniMaxM2Model
+
+
+class DummyWriter:
+    def __init__(self) -> None:
+        self.values: dict[str, object] = {}
+
+    def add_bool(self, key: str, value: bool) -> None:
+        self.values[key] = value
+
+    def add_array(self, key: str, value: list[int]) -> None:
+        self.values[key] = value
+
+
+def make_minimax_model(counts: dict[int, int], *, max_experts: int = 4) -> MiniMaxM2Model:
+    model = object.__new__(MiniMaxM2Model)
+    model.block_count = len(counts)
+    model.hparams = {
+        "hidden_size": 3,
+        "intermediate_size": 2,
+        "num_hidden_layers": len(counts),
+        "num_local_experts": max_experts,
+        "reap_adaptive_nonuniform_experts": True,
+        "reap_adaptive_retained_experts_by_layer": {
+            str(layer): count for layer, count in counts.items()
+        },
+    }
+    model.tensor_map = gguf.get_tensor_name_map(gguf.MODEL_ARCH.MINIMAXM2, model.block_count)
+    model._experts_cache = {}
+    model.fuse_gate_up_exps = False
+    model._adaptive_retained_by_layer = model._load_adaptive_retained_by_layer()
+    model._uniform_expert_count = model.find_hparam(["num_local_experts", "num_experts"])
+    return model
+
+
+def test_minimax_adaptive_export_pads_experts_router_and_bias() -> None:
+    model = make_minimax_model({0: 2, 1: 4})
+
+    outputs: list[tuple[str, torch.Tensor]] = []
+    for expert_id in range(2):
+        for weight_name in ("w1", "w2", "w3"):
+            name = f"model.layers.0.block_sparse_moe.experts.{expert_id}.{weight_name}.weight"
+            tensor = torch.full((2, 3), expert_id + 1.0)
+            outputs.extend(model.modify_tensors(tensor, name, 0))
+
+    assert [name for name, _ in outputs] == [
+        "blk.0.ffn_gate_exps.weight",
+        "blk.0.ffn_down_exps.weight",
+        "blk.0.ffn_up_exps.weight",
+    ]
+    for _, tensor in outputs:
+        assert tensor.shape == (4, 2, 3)
+        assert torch.equal(tensor[2:], torch.zeros_like(tensor[2:]))
+
+    gate_outputs = list(
+        model.modify_tensors(
+            torch.ones((2, 3)),
+            "model.layers.0.block_sparse_moe.gate.weight",
+            0,
+        )
+    )
+    assert len(gate_outputs) == 1
+    assert gate_outputs[0][0] == "blk.0.ffn_gate_inp.weight"
+    assert gate_outputs[0][1].shape == (4, 3)
+    assert torch.equal(gate_outputs[0][1][2:], torch.zeros_like(gate_outputs[0][1][2:]))
+
+    bias_outputs = list(
+        model.modify_tensors(
+            torch.tensor([0.25, 0.5]),
+            "model.layers.0.block_sparse_moe.e_score_correction.bias",
+            0,
+        )
+    )
+    assert len(bias_outputs) == 1
+    assert bias_outputs[0][0] == "blk.0.exp_probs_b.bias"
+    assert bias_outputs[0][1].shape == (4,)
+    # Finite large-negative mask (not -inf): ggml_validate_row_data rejects inf,
+    # and sigmoid-bounded selection scores make -1e9 equally unselectable.
+    assert torch.isfinite(bias_outputs[0][1][2:]).all()
+    assert (bias_outputs[0][1][2:] <= -1.0e8).all()
+
+
+def test_minimax_adaptive_export_records_padding_metadata() -> None:
+    model = make_minimax_model({0: 2, 1: 4})
+    model.gguf_writer = DummyWriter()
+
+    model._add_adaptive_padding_metadata()
+
+    assert model.gguf_writer.values == {
+        "minimax-m2.reap_adaptive_padded_uniform": True,
+        "minimax-m2.reap_adaptive_dummy_experts_masked": True,
+        "minimax-m2.reap_adaptive_retained_experts_by_layer": [2, 4],
+    }
+
+
+def test_minimax_adaptive_export_materializes_lazy_padding_on_source_device() -> None:
+    model = make_minimax_model({0: 2, 1: 4})
+    source = LazyTorchTensor.from_eager(torch.ones((2, 3)))
+
+    padded = model._pad_expert_axis(source, 2, dim=0, fill_value=-1.0)
+    eager = LazyTorchTensor.to_eager(padded)
+
+    assert eager.device.type == "cpu"
+    assert eager.shape == (4, 3)
+    assert torch.equal(eager[:2], torch.ones((2, 3)))
+    assert torch.equal(eager[2:], torch.full((2, 3), -1.0))
+
+
+def test_minimax_adaptive_export_fails_when_metadata_is_incomplete() -> None:
+    model = object.__new__(MiniMaxM2Model)
+    model.block_count = 2
+    model.hparams = {
+        "num_local_experts": 4,
+        "reap_adaptive_nonuniform_experts": True,
+        "reap_adaptive_retained_experts_by_layer": {"0": 2},
+    }
+
+    with pytest.raises(ValueError, match="requires counts for every layer"):
+        model._load_adaptive_retained_by_layer()
+
+
+def test_minimax_adaptive_export_fails_on_incomplete_layer_cache(monkeypatch: pytest.MonkeyPatch) -> None:
+    model = make_minimax_model({0: 2, 1: 4})
+    model._experts_cache = {0: {"partial": torch.ones((2, 3))}}
+    monkeypatch.setattr(TextModel, "prepare_tensors", lambda self: None)
+
+    with pytest.raises(RuntimeError, match="incomplete per-layer expert caches"):
+        model.prepare_tensors()