Fix real PLaMo2 cached decode parity

src/mobius/integrations/gguf/_builder_test.py

@@ -840,6 +840,7 @@ def _write_plamo2_gguf(
     legacy_scalar_heads: bool = False,
     quantized_embedding: bool = False,
     include_output: bool = False,
+    rope_theta: float = 10_000.0,
 ) -> None:
     """Write a complete tiny alternating PLaMo2 GGUF."""
     from gguf import GGMLQuantizationType, GGUFWriter
@@ -867,7 +868,7 @@ def _write_plamo2_gguf(
         kv_heads if legacy_scalar_heads else kv_head_counts or [0, kv_heads]
     )
     writer.add_layer_norm_rms_eps(epsilon)
-    writer.add_rope_freq_base(10_000.0)
+    writer.add_rope_freq_base(rope_theta)
     writer.add_vocab_size(vocab)
     writer.add_ssm_conv_kernel(kernel)
     writer.add_ssm_inner_size(inner)
@@ -4504,7 +4505,7 @@ def _inputs(batch: int) -> dict[str, np.ndarray]:
             "position_ids": np.asarray([[0, 1], [0, 1]][:batch], np.int64),
             "attention_mask": np.ones((batch, 2), np.int64),
             "past_key_values.0.conv_state": np.zeros((batch, 32, 3), np.float32),
-            "past_key_values.0.ssm_state": np.zeros((batch, 4, 8, 4), np.float32),
+            "past_key_values.0.recurrent_state": np.zeros((batch, 4, 8, 4), np.float32),
             "past_key_values.1.key": np.zeros((batch, 2, 0, 8), np.float32),
             "past_key_values.1.value": np.zeros((batch, 2, 0, 8), np.float32),
         }
@@ -4533,13 +4534,13 @@ def test_float_import_executes_and_round_trips(self, tmp_path: Path) -> None:
                 model.graph.initializers[initializer_name].const_value.numpy(),
                 source_tensors[source_name],
             )
-        assert model.metadata_props["mobius.runtime_support"].endswith(
-            "onnxruntime/mobius#605"
+        assert model.metadata_props["mobius.runtime_support"] == (
+            "ORT GenAI 0.15.2 state ABI; package requires GQA-specialized attention"
         )
         assert [value.name for value in model.graph.outputs] == [
             "logits",
             "present.0.conv_state",
-            "present.0.ssm_state",
+            "present.0.recurrent_state",
             "present.1.key",
             "present.1.value",
         ]
@@ -4548,7 +4549,7 @@ def test_float_import_executes_and_round_trips(self, tmp_path: Path) -> None:
         session = OnnxModelSession(ModelPackage.load(output_dir)["model"])
         outputs = session.run(self._inputs(2))
         assert outputs["logits"].shape == (2, 2, 64)
-        assert outputs["present.0.ssm_state"].dtype == np.float32
+        assert outputs["present.0.recurrent_state"].dtype == np.float32
 
     def test_legacy_scalar_heads_infer_exact_tensor_schedule(self, tmp_path: Path) -> None:
         from mobius.integrations.gguf import build_from_gguf
@@ -4560,6 +4561,14 @@ def test_legacy_scalar_heads_infer_exact_tensor_schedule(self, tmp_path: Path) -
         assert package.config.attention_head_counts == (0, 4)
         assert package.config.attention_kv_head_counts == (0, 2)
 
+    def test_legacy_million_base_restores_reference_local_rope(self, tmp_path: Path) -> None:
+        from mobius.integrations.gguf import build_from_gguf
+
+        path = tmp_path / "plamo2-legacy-rope.gguf"
+        _write_plamo2_gguf(path, quantized=False, rope_theta=1_000_000.0)
+        package = build_from_gguf(path)
+        assert package.config.rope_theta == pytest.approx(10_000.0)
+
     def test_quantized_source_preserves_exact_projection_roles(self, tmp_path: Path) -> None:
         from mobius.integrations.gguf import build_from_gguf
 
@@ -4568,7 +4577,7 @@ def test_quantized_source_preserves_exact_projection_roles(self, tmp_path: Path)
         model = build_from_gguf(path, keep_quantized=True)["model"]
         assert sum(node.op_type == "MatMulNBits" for node in model.graph) == 9
         assert (
-            model.graph.initializers["model.layers.0.mixer.A"].const_value.dtype
+            model.graph.initializers["model.layers.0.mixer.A_log"].const_value.dtype
             == ir.DataType.FLOAT
         )
         assert (

src/mobius/integrations/gguf/_config_mapping.py

@@ -2244,6 +2244,14 @@ def _plamo2_postprocess(
     fields.update(
         hidden_act="silu",
         head_dim=key_length,
+        # The released PLaMo2 converter wrote 1e6 here even though the pinned
+        # reference architecture uses its 1e4 local-RoPE default for every
+        # attention layer. Preserve other explicit bases for future variants.
+        rope_theta=(
+            10_000.0
+            if float(metadata[f"{arch}.rope.freq_base"]) == 1_000_000.0  # noqa: RUF069
+            else config.rope_theta
+        ),
         attention_head_counts=head_counts,
         attention_kv_head_counts=kv_head_counts,
         mamba_num_heads=ssm_heads,

src/mobius/integrations/ort_genai/auto_export.py

@@ -108,6 +108,9 @@ def _revision_kwargs(revision: str | None) -> dict[str, str]:
     # ORT GenAI (see onnxruntime-genai/src/models/model_type.h LLM list).
     "hunyuan_v1_dense": "decoder",
     "deepseek_v4": "decoder",
+    # PLaMo2 is a decoder-only hybrid. Released ORT GenAI does not have a
+    # model-specific registry entry, so emit its generic decoder type.
+    "plamo2": "decoder",
     # Qwen VL model families have separate ORT GenAI model types.
     "qwen2_vl": "qwen2_5_vl",
     "qwen3_vl": "qwen3_vl",
@@ -202,13 +205,15 @@ def _revision_kwargs(revision: str | None) -> dict[str, str]:
 
 _TOKENIZER_FILES = [
     "tokenizer.json",
+    "tokenizer.jsonl",  # PLaMo2 scored vocabulary
     "tokenizer_config.json",
     "special_tokens_map.json",
     "tokenizer.model",  # SentencePiece
     "added_tokens.json",
     "merges.txt",  # BPE
     "vocab.json",  # BPE
     "chat_template.jinja",  # Chat template for ORT GenAI
+    "tokenization_plamo.py",  # PLaMo2's exact custom tokenizer implementation
     # Preserve HuggingFace processor metadata for VLMs whose preprocessing
     # cannot be represented by an ort-extensions image_processor.json.
     "preprocessor_config.json",
@@ -477,6 +482,7 @@ def _fix_chat_template(
     hf_model_id: str | None,
     *,
     revision: str | None = None,
+    trust_remote_code: bool = False,
 ) -> bool:
     """Ensure chat_template is present in tokenizer_config.json.
 
@@ -508,6 +514,7 @@ def _fix_chat_template(
 
         tokenizer = AutoTokenizer.from_pretrained(
             hf_model_id,
+            trust_remote_code=trust_remote_code,
             **_revision_kwargs(revision),
         )
         template = getattr(tokenizer, "chat_template", None)
@@ -1760,7 +1767,12 @@ def write_ort_genai_config(
     _fix_tokenizer_config(directory)
 
     # Ensure chat_template is in tokenizer_config.json
-    _fix_chat_template(directory, hf_model_id, revision=revision)
+    _fix_chat_template(
+        directory,
+        hf_model_id,
+        revision=revision,
+        trust_remote_code=trust_remote_code,
+    )
 
     # Correct assets that ship broken from upstream
     apply_asset_patches(directory)

src/mobius/integrations/ort_genai/auto_export_test.py

@@ -110,6 +110,9 @@ def test_hunyuan_v1_dense_maps_to_decoder(self):
         # decoder-only causal LM not in its built-in registry.
         assert _resolve_ort_genai_model_type("hunyuan_v1_dense") == "decoder"
 
+    def test_plamo2_maps_to_generic_decoder(self):
+        assert _resolve_ort_genai_model_type("plamo2") == "decoder"
+
     def test_unknown_model_type_passthrough(self):
         assert _resolve_ort_genai_model_type("my_custom") == "my_custom"
 
@@ -807,10 +810,20 @@ def test_adds_chat_template(self, tmp_path):
         with mock.patch(
             "transformers.AutoTokenizer.from_pretrained",
             return_value=fake_tokenizer,
-        ):
-            result = _fix_chat_template(str(tmp_path), "fake/model")
+        ) as from_pretrained:
+            result = _fix_chat_template(
+                str(tmp_path),
+                "fake/model",
+                revision="immutable-revision",
+                trust_remote_code=True,
+            )
 
         assert result is True
+        from_pretrained.assert_called_once_with(
+            "fake/model",
+            revision="immutable-revision",
+            trust_remote_code=True,
+        )
         fixed = json.loads((tmp_path / "tokenizer_config.json").read_text())
         assert fixed["chat_template"] == "{{ bos_token }}"
 
@@ -946,7 +959,9 @@ def test_copies_present_files(self, tmp_path):
         src = tmp_path / "model"
         src.mkdir()
         (src / "tokenizer.json").write_text('{"test": true}')
+        (src / "tokenizer.jsonl").write_text('["token", 0.0, "NORMAL"]\n')
         (src / "tokenizer_config.json").write_text('{"model_type": "llama"}')
+        (src / "tokenization_plamo.py").write_text("class Plamo2Tokenizer: pass\n")
         (src / "chat_template.jinja").write_text("{{ messages }}")
 
         dst = tmp_path / "output"
@@ -955,7 +970,9 @@ def test_copies_present_files(self, tmp_path):
 
         assert set(copied) == {
             "tokenizer.json",
+            "tokenizer.jsonl",
             "tokenizer_config.json",
+            "tokenization_plamo.py",
             "chat_template.jinja",
         }
         assert (dst / "tokenizer.json").read_text() == '{"test": true}'

src/mobius/models/plamo2.py

@@ -140,8 +140,7 @@ def __init__(self, config: Plamo2Config):
         # dt participates in Softplus and remains float even for quantized GGUF imports.
         self.dt_proj = Linear(self.dt_rank, self.num_heads, bias=False)
         self.dt_bias = nn.Parameter([self.num_heads])
-        # GGUF stores the already transformed negative decay directly.
-        self.A = nn.Parameter([self.num_heads])
+        self.A_log = nn.Parameter([self.num_heads])
         self.D = nn.Parameter([self.num_heads])
         self.dt_norm_weight = nn.Parameter([self.dt_rank])
         self.B_norm_weight = nn.Parameter([self.state_size])
@@ -198,7 +197,10 @@ def forward(
             )
         )
         dt = op.Mul(dt, op.CastLike(padding_mask, dt))
-        decay = op.Mul(dt, op.Cast(self.A, to=ir.DataType.FLOAT))
+        # Match the reference's float32 discretization at runtime rather than
+        # baking a rounded -exp(A_log) into imported weights.
+        a = op.Neg(op.Exp(op.Cast(self.A_log, to=ir.DataType.FLOAT)))
+        decay = op.Mul(dt, a)
         x_f32 = op.Cast(x, to=ir.DataType.FLOAT)
         x_heads = op.Reshape(x_f32, [0, 0, self.num_heads, self.head_dim])
         value = op.Reshape(
@@ -414,7 +416,7 @@ def preprocess_weights(
         self,
         state_dict: dict[str, torch.Tensor],
     ) -> dict[str, torch.Tensor]:
-        """Convert the official offset norms and Mamba decay to graph values."""
+        """Convert official and GGUF offset norms and Mamba decay names."""
         result: dict[str, torch.Tensor] = {}
         norm_offsets = {
             ".pre_mixer_norm.weight": 1.0,
@@ -423,9 +425,16 @@ def preprocess_weights(
             ".post_mlp_norm.weight": 1.0 / (5.0**1.5),
         }
         norms_are_folded = bool(getattr(self.config, "_plamo2_norms_are_folded", False))
+        tied_embeddings = effective_tie_word_embeddings(self.config)
         for name, value in state_dict.items():
             name = name.replace("model.layers.layers.", "model.layers.")
-            if name == "lm_head.weight" and self.config.tie_word_embeddings:
+            if name == "model.embed_tokens.weight" and tied_embeddings:
+                result[name] = value
+                # onnxscript materializes the shared Parameter under both use
+                # sites, while the official checkpoint stores only the embedding.
+                result["lm_head.weight"] = value
+                continue
+            if name == "lm_head.weight" and tied_embeddings:
                 continue
             if name == "model.norm.weight":
                 result[name] = value if norms_are_folded else value + 1.0
@@ -436,8 +445,9 @@ def preprocess_weights(
             )
             if offset is not None:
                 result[name] = value if norms_are_folded else value + offset
-            elif name.endswith(".mixer.A_log"):
-                result[name.removesuffix("A_log") + "A"] = -torch.exp(value)
+            elif name.endswith(".mixer.A"):
+                # llama.cpp serializes PLaMo2's A_log as -exp(A_log).
+                result[name.removesuffix("A") + "A_log"] = torch.log(-value)
             else:
                 result[name] = value
         return result

src/mobius/models/plamo2_test.py

@@ -10,7 +10,7 @@
 import torch
 
 from mobius import build_from_module
-from mobius._configs import Plamo2Config
+from mobius._configs import Plamo2Config, QuantizationConfig
 from mobius._testing.ort_inference import OnnxModelSession
 from mobius.models.plamo2 import Plamo2ForCausalLM
 from mobius.tasks import Plamo2CausalLMTask
@@ -50,9 +50,7 @@ def _fill_weights(model: ir.Model) -> None:
             continue
         shape = [dim if isinstance(dim, int) else 1 for dim in initializer.shape]
         values = (rng.standard_normal(shape) * 0.03).astype(initializer.dtype.numpy())
-        if initializer.name.endswith(".A"):
-            values = -np.exp(values)
-        elif "norm" in initializer.name:
+        if "norm" in initializer.name:
             values += 1.0
         initializer.const_value = ir.Tensor(values)
 
@@ -62,7 +60,7 @@ def _empty_states(config: Plamo2Config, batch: int) -> dict[str, np.ndarray]:
         "past_key_values.0.conv_state": np.zeros(
             (batch, config.mamba_inner_size, config.mamba_d_conv - 1), np.float32
         ),
-        "past_key_values.0.ssm_state": np.zeros(
+        "past_key_values.0.recurrent_state": np.zeros(
             (
                 batch,
                 config.mamba_num_heads,
@@ -124,10 +122,8 @@ def test_plamo2_mamba_matches_independent_reduced_reference() -> None:
         initializer.const_value = ir.Tensor(value)
         values[initializer.name] = value
 
-    def assign(name: str, scale: float = 0.05, *, negative: bool = False) -> np.ndarray:
+    def assign(name: str, scale: float = 0.05) -> np.ndarray:
         value = (rng.standard_normal(values[name].shape) * scale).astype(np.float32)
-        if negative:
-            value = -np.exp(value)
         model.graph.initializers[name].const_value = ir.Tensor(value)
         values[name] = value
         return value
@@ -138,7 +134,7 @@ def assign(name: str, scale: float = 0.05, *, negative: bool = False) -> np.ndar
     assign("model.layers.0.mixer.bcdt_proj.weight")
     assign("model.layers.0.mixer.dt_proj.weight")
     assign("model.layers.0.mixer.dt_bias")
-    assign("model.layers.0.mixer.A", negative=True)
+    assign("model.layers.0.mixer.A_log")
     assign("model.layers.0.mixer.D")
     assign("model.layers.0.mixer.out_proj.weight")
 
@@ -204,7 +200,8 @@ def assign(name: str, scale: float = 0.05, *, negative: bool = False) -> np.ndar
     scan_outputs = []
     for index in range(input_ids.shape[1]):
         decay = np.exp(
-            dt[:, index, :, None, None] * values["model.layers.0.mixer.A"][None, :, None, None]
+            dt[:, index, :, None, None]
+            * -np.exp(values["model.layers.0.mixer.A_log"])[None, :, None, None]
         )
         update = (
             dt[:, index, :, None, None]
@@ -229,7 +226,9 @@ def assign(name: str, scale: float = 0.05, *, negative: bool = False) -> np.ndar
         rtol=1e-6,
         atol=1e-7,
     )
-    np.testing.assert_allclose(outputs["present.0.ssm_state"], state, rtol=2e-5, atol=2e-6)
+    np.testing.assert_allclose(
+        outputs["present.0.recurrent_state"], state, rtol=2e-5, atol=2e-6
+    )
 
 
 def test_plamo2_attention_and_mlp_match_independent_reduced_reference() -> None:
@@ -407,7 +406,9 @@ def test_plamo2_left_padding_does_not_change_recurrent_state_or_valid_logits() -
         padded["present.0.conv_state"], unpadded["present.0.conv_state"], atol=1e-7
     )
     np.testing.assert_allclose(
-        padded["present.0.ssm_state"], unpadded["present.0.ssm_state"], atol=1e-7
+        padded["present.0.recurrent_state"],
+        unpadded["present.0.recurrent_state"],
+        atol=1e-7,
     )
 
 
@@ -451,6 +452,7 @@ def test_plamo2_preprocesses_offsets_and_decay_values() -> None:
         "model.layers.layers.0.pre_mlp_norm.weight": torch.tensor([5.0]),
         "model.layers.layers.0.post_mlp_norm.weight": torch.tensor([6.0]),
         "model.layers.layers.0.mixer.A_log": torch.tensor([0.0, 1.0]),
+        "model.embed_tokens.weight": torch.tensor([8.0]),
         "lm_head.weight": torch.tensor([9.0]),
     }
     actual = model.preprocess_weights(weights)
@@ -468,10 +470,31 @@ def test_plamo2_preprocesses_offsets_and_decay_values() -> None:
         actual["model.layers.0.post_mlp_norm.weight"],
         torch.tensor([6.0 + 1.0 / (5.0**1.5)]),
     )
+    torch.testing.assert_close(actual["model.layers.0.mixer.A_log"], torch.tensor([0.0, 1.0]))
+    gguf_actual = model.preprocess_weights(
+        {"model.layers.0.mixer.A": -torch.exp(torch.tensor([0.0, 1.0]))}
+    )
     torch.testing.assert_close(
-        actual["model.layers.0.mixer.A"], -torch.exp(torch.tensor([0.0, 1.0]))
+        gguf_actual["model.layers.0.mixer.A_log"], torch.tensor([0.0, 1.0])
+    )
+    torch.testing.assert_close(actual["lm_head.weight"], torch.tensor([8.0]))
+
+
+def test_plamo2_preprocesses_effectively_tied_quantized_embeddings() -> None:
+    config = _config()
+    config.tie_word_embeddings = False
+    config.quantization = QuantizationConfig(tie_word_embeddings=True)
+    model = Plamo2ForCausalLM(config)
+
+    actual = model.preprocess_weights(
+        {
+            "model.embed_tokens.weight": torch.tensor([8.0]),
+            "lm_head.weight": torch.tensor([9.0]),
+        }
     )
-    assert "lm_head.weight" not in actual
+
+    torch.testing.assert_close(actual["model.embed_tokens.weight"], torch.tensor([8.0]))
+    torch.testing.assert_close(actual["lm_head.weight"], torch.tensor([8.0]))
 
 
 def test_plamo2_transformers_config_uses_explicit_schedule_and_pinned_defaults() -> None: