Add HRM-Text (hrm_text) hierarchical recurrent model with PrefixLM masking

scripts/generate_golden.py

@@ -244,6 +244,10 @@ def _get_model_device(model: object, device: str):
 def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
     """Generate golden data for a causal-lm (text-generation) model."""
     from mobius._testing.golden import save_generation_json, save_golden_ref
+    from mobius._testing.prefix_lm import (
+        model_type_uses_prefix_lm,
+        prompt_token_type_ids,
+    )
     from mobius._testing.torch_reference import (
         load_torch_model,
         load_torch_multimodal_model,
@@ -267,7 +271,10 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
         )
     else:
         model, tokenizer = load_torch_model(
-            case.model_id, device=device, trust_remote_code=case.trust_remote_code
+            case.model_id,
+            device=device,
+            trust_remote_code=case.trust_remote_code,
+            revision=case.revision,
         )
 
     encoded = tokenizer(case.prompts[0], return_tensors="np", padding=False)
@@ -276,6 +283,13 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
     seq_len = input_ids.shape[1]
     position_ids = np.arange(seq_len).reshape(1, -1)
 
+    # PrefixLM models (HRM-Text) must run the reference with the same
+    # ``token_type_ids`` contract the exported graph is fed: the whole prompt is
+    # one bidirectional prefix block. Omitting it silently falls back to causal
+    # masking, which the model card warns does not match pre-training.
+    uses_prefix_lm = model_type_uses_prefix_lm(case.model_type)
+    prefix_token_type_ids = prompt_token_type_ids(input_ids) if uses_prefix_lm else None
+
     # L4: single forward pass → last-token logits
     if uses_multimodal_reference:
         with torch.no_grad():
@@ -287,7 +301,13 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
             )
         logits = outputs.logits.float().cpu().numpy()
     else:
-        logits, _ = torch_forward(model, input_ids, attention_mask, position_ids)
+        logits, _ = torch_forward(
+            model,
+            input_ids,
+            attention_mask,
+            position_ids,
+            token_type_ids=prefix_token_type_ids,
+        )
     last_logits = logits[0, -1, :]  # (vocab_size,)
     golden = _extract_logits_golden(last_logits)
 
@@ -301,9 +321,19 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
         model_device = _get_model_device(model, device)
         gen_ids = torch.from_numpy(input_ids).to(model_device)
         max_new = case.generation_params.get("max_new_tokens", 20)
+        gen_kwargs: dict = {"max_new_tokens": max_new, "do_sample": False}
+        if prefix_token_type_ids is not None:
+            # HF's HrmTextForCausalLM.create_masks_for_generate applies the
+            # PrefixLM overlay only on the first iteration, so marking the whole
+            # prompt as one prefix block is exactly the model card's
+            # ``token_type_ids = ones_like(input_ids)`` call.
+            gen_kwargs["attention_mask"] = torch.from_numpy(attention_mask).to(model_device)
+            gen_kwargs["token_type_ids"] = torch.from_numpy(prefix_token_type_ids).to(
+                model_device
+            )
         with torch.no_grad():
             try:
-                gen_output = model.generate(gen_ids, max_new_tokens=max_new, do_sample=False)
+                gen_output = model.generate(gen_ids, **gen_kwargs)
             except ValueError as e:
                 # All-attention GraniteMoeHybrid variants (e.g. granite-4.0-1b)
                 # trip transformers' hybrid Mamba/attention generation cache,
@@ -313,9 +343,7 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
                 # the (slower) cache-free path.
                 if "has_previous_state" not in str(e):
                     raise
-                gen_output = model.generate(
-                    gen_ids, max_new_tokens=max_new, do_sample=False, use_cache=False
-                )
+                gen_output = model.generate(gen_ids, use_cache=False, **gen_kwargs)
         generated_ids = gen_output[0, seq_len:].cpu().numpy()
 
     save_golden_ref(

src/mobius/_configs/__init__.py

@@ -38,6 +38,7 @@
     Gemma4Config,
     GlmAsrConfig,
     GraniteMoeHybridConfig,
+    HrmTextConfig,
     JambaConfig,
     JetMoeConfig,
     Lfm2Config,
@@ -114,6 +115,7 @@
     "Gemma4Config",
     "GlmAsrConfig",
     "GraniteMoeHybridConfig",
+    "HrmTextConfig",
     "JambaConfig",
     "JetMoeConfig",
     "Lfm2Config",

src/mobius/_configs/_base.py

@@ -1739,6 +1739,109 @@ def from_transformers(cls, config, parent_config=None) -> Gemma2Config:
         )
 
 
+def _as_int_or_default(value, default: int) -> int:
+    """Return ``int(value)``, falling back to *default* only when unset.
+
+    Unlike ``int(value or default)`` this preserves an explicit ``0`` so that
+    an out-of-range checkpoint value reaches the config's own validation
+    instead of being silently rewritten to the default.
+    """
+    return default if value is None else int(value)
+
+
+def _as_float_or_default(value, default: float) -> float:
+    """Float counterpart of :func:`_as_int_or_default`."""
+    return default if value is None else float(value)
+
+
+@dataclasses.dataclass
+class HrmTextConfig(CausalLMConfig):
+    """Configuration for HRM-Text hierarchical recurrent models.
+
+    Mirrors HuggingFace ``HrmTextConfig``. The checkpoint stores the real
+    per-stack block count in ``num_hidden_layers``; HuggingFace's
+    ``__post_init__`` moves it to ``num_layers_per_stack`` and rewrites
+    ``num_hidden_layers`` to the *inflated* total number of unique attention
+    invocations under the H/L recurrence::
+
+        num_hidden_layers = num_layers_per_stack * H_cycles * (L_cycles + 1)
+
+    That inflated count is what drives KV-cache slot allocation, so it is what
+    :class:`~mobius.tasks.CausalLMTask` must see. This dataclass reproduces the
+    same split so that both a trusted ``HrmTextConfig`` instance (already
+    inflated, ``num_layers_per_stack`` set) and a raw pinned ``config.json``
+    (not inflated, ``num_layers_per_stack`` absent) resolve identically.
+
+    ``embedding_scale`` defaults to ``1 / initializer_range`` exactly as
+    upstream does when the checkpoint leaves it unset.
+    """
+
+    H_cycles: int = 2
+    L_cycles: int = 3
+    num_layers_per_stack: int | None = None
+    embedding_scale: float | None = None
+    initializer_range: float = 0.02
+    prefix_lm: bool = True
+
+    def __post_init__(self):
+        if self.H_cycles <= 0 or self.L_cycles <= 0:
+            raise ValueError(
+                f"HrmTextConfig requires positive H_cycles/L_cycles, got "
+                f"H_cycles={self.H_cycles}, L_cycles={self.L_cycles}"
+            )
+        # HRM-Text attention is always MHA: upstream hardcodes
+        # ``num_key_value_groups = 1`` and sizes k_proj/v_proj by
+        # ``num_attention_heads * head_dim``, ignoring any
+        # ``num_key_value_heads`` the checkpoint happens to carry.
+        if self.num_attention_heads != DEFAULT_INT:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.embedding_scale is None:
+            if not self.initializer_range:
+                raise ValueError(
+                    "HrmTextConfig needs a non-zero initializer_range to derive "
+                    "embedding_scale when the checkpoint does not supply one."
+                )
+            self.embedding_scale = 1.0 / self.initializer_range
+        if self.num_layers_per_stack is None and self.num_hidden_layers != DEFAULT_INT:
+            # Raw-config path: ``num_hidden_layers`` still carries the real
+            # per-stack count. Remember it, then inflate exactly as upstream.
+            self.num_layers_per_stack = self.num_hidden_layers
+            self.num_hidden_layers = (
+                self.num_layers_per_stack * self.H_cycles * (self.L_cycles + 1)
+            )
+
+    @classmethod
+    def from_transformers(cls, config, parent_config=None) -> HrmTextConfig:
+        base = ArchitectureConfig.from_transformers(config, parent_config)
+        fields = _shallow_fields(base)
+        # ``HrmTextRotaryEmbedding`` is unconditional upstream, so this family
+        # always uses RoPE. A raw ``config.json`` only carries the default
+        # ``rope_theta`` of 10000.0, which the generic extractor deliberately
+        # ignores as a RoPE signal (NoPE models inherit it as dead config
+        # data), leaving ``rope_type=None`` and silently exporting a
+        # position-free graph. Pin the default RoPE explicitly instead.
+        if fields.get("rope_type") is None:
+            fields["rope_type"] = "default"
+            if fields.get("rope_theta") is None:
+                fields["rope_theta"] = 10_000.0
+            if fields.get("partial_rotary_factor") is None:
+                # Upstream rotates the full head_dim.
+                fields["partial_rotary_factor"] = 1.0
+        return cls(
+            **fields,
+            H_cycles=_as_int_or_default(getattr(config, "H_cycles", None), 2),
+            L_cycles=_as_int_or_default(getattr(config, "L_cycles", None), 3),
+            # Present only on a trusted HF config that already ran its own
+            # ``__post_init__``; ``None`` for a raw ``config.json``.
+            num_layers_per_stack=getattr(config, "num_layers_per_stack", None),
+            embedding_scale=getattr(config, "embedding_scale", None),
+            initializer_range=_as_float_or_default(
+                getattr(config, "initializer_range", None), 0.02
+            ),
+            prefix_lm=bool(getattr(config, "prefix_lm", True)),
+        )
+
+
 @dataclasses.dataclass
 class NanoChatConfig(CausalLMConfig):
     """Configuration for NanoChat models with final logit soft-capping.

src/mobius/_registry.py

@@ -28,6 +28,7 @@
     Gemma3nMultiModalConfig,
     Gemma4AssistantConfig,
     Gemma4Config,
+    HrmTextConfig,
     Lfm2Config,
     Lfm2VlConfig,
     MMSConfig,
@@ -73,6 +74,7 @@
     GPTOSSCausalLMModel,
     GraniteCausalLMModel,
     GraniteMoECausalLMModel,
+    HrmTextCausalLMModel,
     HunYuanMoEV1CausalLMModel,
     HunYuanV1DenseCausalLMModel,
     HunYuanVLMoTModel,
@@ -492,6 +494,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "gpt_oss": ModelRegistration(GPTOSSCausalLMModel),
     "gptj": ModelRegistration(GPTJCausalLMModel),
     "granite": ModelRegistration(GraniteCausalLMModel),
+    "hrm_text": ModelRegistration(HrmTextCausalLMModel, config_class=HrmTextConfig),
     "hunyuan_v1_dense": ModelRegistration(HunYuanV1DenseCausalLMModel),
     "internlm2": ModelRegistration(InternLM2CausalLMModel),
     "llama4_text": ModelRegistration(Llama4CausalLMModel),
@@ -953,6 +956,7 @@ def _create_default_registry() -> ModelRegistry:
     "dots1": "rednote-hilab/dots.llm1.inst",
     "exaone4": "LGAI-EXAONE/EXAONE-4.0-1.2B",
     "helium": "kyutai/helium-1-preview-2b",
+    "hrm_text": "sapientinc/HRM-Text-1B",
     "minicpm": "optimum-intel-internal-testing/tiny-random-minicpm",
     "minicpm3": "openbmb/MiniCPM3-4B",
     "ministral3": "Aratako/Ministral-3-3B-Instruct-2512-BF16-TextOnly",

src/mobius/_testing/generation.py

@@ -15,6 +15,11 @@
 
 from mobius._configs import ArchitectureConfig
 from mobius._testing.ort_inference import OnnxModelSession
+from mobius._testing.prefix_lm import (
+    TOKEN_TYPE_IDS,
+    generated_token_type_ids,
+    prompt_token_type_ids,
+)
 
 
 class OnnxGenerator:
@@ -109,6 +114,12 @@ def _kv_dtype(name: str) -> np.dtype:
 
         all_ids = input_ids.copy()
 
+        # PrefixLM models (HRM-Text) take an extra ``token_type_ids`` input:
+        # the whole prompt is one bidirectional prefix block (1), every
+        # generated token is causal (0). Mirrors the model card's documented
+        # ``token_type_ids = ones_like(input_ids)`` prefill call.
+        uses_token_type_ids = TOKEN_TYPE_IDS in set(self.session.input_names)
+
         # First step: process the full prompt
         cur_input_ids = input_ids
         past_seq_len = 0
@@ -128,6 +139,12 @@ def _kv_dtype(name: str) -> np.dtype:
                 "position_ids": position_ids,
                 **past_kv,
             }
+            if uses_token_type_ids:
+                feeds[TOKEN_TYPE_IDS] = (
+                    prompt_token_type_ids(cur_input_ids)
+                    if past_seq_len == 0
+                    else generated_token_type_ids(cur_input_ids)
+                )
 
             outputs = self.session.run(feeds)
 

src/mobius/_testing/prefix_lm.py

@@ -0,0 +1,54 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""PrefixLM ``token_type_ids`` contract shared by every test/golden path.
+
+Models whose exported graph declares ``requires_token_type_ids`` were
+pre-trained with a PrefixLM mask: prompt tokens attend bidirectionally to each
+other, generated tokens attend causally.  ``sapientinc/HRM-Text-1B``'s model
+card spells the correct call out explicitly — mark the *entire* prompt as one
+bidirectional prefix block (``token_type_ids = ones_like(input_ids)``) and
+leave generated positions at ``0``.
+
+Both sides of every comparison must use the same rule or the numbers are
+meaningless, so the rule lives here exactly once: the HuggingFace reference is
+driven from :func:`model_type_uses_prefix_lm` and the exported graph from
+:func:`graph_uses_prefix_lm`, and both hand out the same tensors.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    import onnx_ir as ir
+
+TOKEN_TYPE_IDS = "token_type_ids"
+
+
+def model_type_uses_prefix_lm(model_type: str) -> bool:
+    """Whether *model_type*'s registered class consumes ``token_type_ids``."""
+    from mobius._registry import registry
+
+    try:
+        model_cls = registry.get(model_type)
+    except Exception:
+        return False
+    return bool(getattr(model_cls, "requires_token_type_ids", False))
+
+
+def graph_uses_prefix_lm(model: ir.Model) -> bool:
+    """Whether an exported graph declares a ``token_type_ids`` input."""
+    return any(value.name == TOKEN_TYPE_IDS for value in model.graph.inputs)
+
+
+def prompt_token_type_ids(input_ids: np.ndarray) -> np.ndarray:
+    """``token_type_ids`` for a prompt that is one bidirectional prefix block."""
+    return np.ones_like(input_ids, dtype=np.int64)
+
+
+def generated_token_type_ids(input_ids: np.ndarray) -> np.ndarray:
+    """``token_type_ids`` for generated positions: causal, never prefix."""
+    return np.zeros_like(input_ids, dtype=np.int64)