hrm_mlx.py

"""MLX implementation of HRM-Text (Sapient Intelligence Hierarchical Reasoning Model).

Faithful port of `sapientinc/HRM-Text-1B` with:
  - 128-slot recurrent KV cache (one per cycle-step × layer)
  - mx.fast.scaled_dot_product_attention (Metal-optimized)
  - mx.fast.rope with cache-aware position offset
  - mx.fast.rms_norm (parameterless, weight=None)
  - Greedy + temperature/top-p sampling
  - Streaming generation
  - Prompt formatting with HRM condition modes

Architecture (per the paper / HF model card):
  z_H = embed(input_ids) * embedding_scale     # slow / strategic
  z_L = zeros                                  # fast / execution
  for h in range(H_cycles):
      for l in range(L_cycles):
          z_L = L_module(z_L + z_H)            # shared transformer stack
      z_H = H_module(z_H + z_L)                # shared transformer stack
  logits = lm_head(z_H)

Each stack invocation maps to its own KV cache slot:
  slot(h, l, layer) = (h * (L_cycles + 1) + l) * num_layers_per_stack + layer    # L-module call (h, l)
  slot(h, layer)    = (h * (L_cycles + 1) + L_cycles) * num_layers_per_stack + layer  # H-module call (h)

Weight loading reads the fused HF safetensors layout 1:1:
  attn.gqkv_proj  [4H, H]   — chunked on dim 0 as [gate, q, k, v]
  mlp.gate_up_proj [2I, H]  — chunked on dim 0 as [gate, up]
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator, Optional, Union

import mlx.core as mx
import mlx.core.fast as mxf
import mlx.nn as nn
from safetensors import safe_open


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
@dataclass
class HrmTextConfig:
    vocab_size: int
    hidden_size: int
    intermediate_size: int
    num_layers_per_stack: int
    num_attention_heads: int
    head_dim: int
    max_position_embeddings: int
    rms_norm_eps: float
    rope_theta: float
    H_cycles: int
    L_cycles: int
    embedding_scale: float
    prefix_lm: bool = True
    tie_word_embeddings: bool = False
    pad_token_id: Optional[int] = None
    bos_token_id: Optional[int] = None
    eos_token_id: Optional[int] = None

    @classmethod
    def from_dict(cls, d: dict) -> "HrmTextConfig":
        return cls(
            vocab_size=d["vocab_size"],
            hidden_size=d["hidden_size"],
            intermediate_size=d["intermediate_size"],
            num_layers_per_stack=d["num_hidden_layers"],  # checkpoint stores per-stack count
            num_attention_heads=d["num_attention_heads"],
            head_dim=d["head_dim"],
            max_position_embeddings=d["max_position_embeddings"],
            rms_norm_eps=d["rms_norm_eps"],
            rope_theta=d["rope_theta"],
            H_cycles=d["H_cycles"],
            L_cycles=d["L_cycles"],
            embedding_scale=d["embedding_scale"],
            prefix_lm=d.get("prefix_lm", True),
            tie_word_embeddings=d.get("tie_word_embeddings", False),
            pad_token_id=d.get("pad_token_id"),
            bos_token_id=d.get("bos_token_id"),
            eos_token_id=d.get("eos_token_id"),
        )

    @property
    def num_cache_slots(self) -> int:
        return self.H_cycles * (self.L_cycles + 1) * self.num_layers_per_stack


# ---------------------------------------------------------------------------
# KV cache
# ---------------------------------------------------------------------------
class _KVSlot:
    """Single layer/cycle KV slot with chunked-grow allocation (mlx-lm style)."""

    __slots__ = ("keys", "values", "offset", "step")

    def __init__(self, step: int = 256):
        self.keys: Optional[mx.array] = None
        self.values: Optional[mx.array] = None
        self.offset: int = 0
        self.step: int = step

    def update_and_fetch(self, k_new: mx.array, v_new: mx.array) -> tuple[mx.array, mx.array]:
        """Append ``k_new``/``v_new`` to the slot and return the cumulative views.

        ``k_new``, ``v_new`` shape: (B, H, T_new, D).
        """
        prev = self.offset
        new_len = prev + k_new.shape[2]

        if self.keys is None or new_len > self.keys.shape[2]:
            B, H, _, D = k_new.shape
            grow_to = self.step * ((new_len + self.step - 1) // self.step)
            grown_k = mx.zeros((B, H, grow_to, D), dtype=k_new.dtype)
            grown_v = mx.zeros((B, H, grow_to, D), dtype=v_new.dtype)
            if self.keys is not None:
                grown_k[..., :prev, :] = self.keys[..., :prev, :]
                grown_v[..., :prev, :] = self.values[..., :prev, :]
            self.keys = grown_k
            self.values = grown_v

        self.keys[..., prev:new_len, :] = k_new
        self.values[..., prev:new_len, :] = v_new
        self.offset = new_len
        return self.keys[..., :new_len, :], self.values[..., :new_len, :]


class HrmKVCache:
    """128-slot recurrent cache (one per cycle-step × layer)."""

    def __init__(self, cfg: HrmTextConfig, step: int = 256):
        self.cfg = cfg
        self.slots = [_KVSlot(step) for _ in range(cfg.num_cache_slots)]

    @property
    def offset(self) -> int:
        """Number of sequence positions cached (each slot advances in lockstep)."""
        return self.slots[0].offset

    def slot(self, idx: int) -> _KVSlot:
        return self.slots[idx]


# ---------------------------------------------------------------------------
# Building blocks
# ---------------------------------------------------------------------------
class HrmAttention(nn.Module):
    """MHA + sigmoid output gate + RoPE. Fused gqkv projection: [gate | q | k | v]."""

    def __init__(self, hidden_size: int, num_heads: int, head_dim: int, rope_theta: float):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = 1.0 / math.sqrt(head_dim)
        self.rope_theta = rope_theta
        self.gqkv_proj = nn.Linear(hidden_size, 4 * num_heads * head_dim, bias=False)
        self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[Union[str, mx.array]],
        rope_offset: int,
        slot: Optional[_KVSlot] = None,
    ) -> mx.array:
        B, S, _ = x.shape
        H, D = self.num_heads, self.head_dim

        gqkv = self.gqkv_proj(x)
        gate, q, k, v = mx.split(gqkv, 4, axis=-1)
        q = q.reshape(B, S, H, D).transpose(0, 2, 1, 3)
        k = k.reshape(B, S, H, D).transpose(0, 2, 1, 3)
        v = v.reshape(B, S, H, D).transpose(0, 2, 1, 3)
        gate = gate.reshape(B, S, H, D)

        # RoPE applied to q, k only — Llama / GPT-NeoX rotate_half convention.
        q = mxf.rope(q, D, traditional=False, base=self.rope_theta, scale=1.0, offset=rope_offset)
        k = mxf.rope(k, D, traditional=False, base=self.rope_theta, scale=1.0, offset=rope_offset)

        if slot is not None:
            k, v = slot.update_and_fetch(k, v)

        out = mxf.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
        out = out.transpose(0, 2, 1, 3)              # (B, S, H, D)
        out = mx.sigmoid(gate) * out                  # output gating
        out = out.reshape(B, S, H * D)
        return self.o_proj(out)


class HrmMLP(nn.Module):
    """SwiGLU with fused gate_up projection: [gate | up]."""

    def __init__(self, hidden_size: int, intermediate_size: int):
        super().__init__()
        self.gate_up_proj = nn.Linear(hidden_size, 2 * intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def __call__(self, x: mx.array) -> mx.array:
        gu = self.gate_up_proj(x)
        gate, up = mx.split(gu, 2, axis=-1)
        return self.down_proj(nn.silu(gate) * up)


class HrmLayer(nn.Module):
    """Pre-norm decoder block with parameterless RMSNorm."""

    def __init__(self, cfg: HrmTextConfig):
        super().__init__()
        self.attn = HrmAttention(cfg.hidden_size, cfg.num_attention_heads, cfg.head_dim, cfg.rope_theta)
        self.mlp = HrmMLP(cfg.hidden_size, cfg.intermediate_size)
        self.eps = cfg.rms_norm_eps

    def __call__(
        self,
        x: mx.array,
        mask: Optional[Union[str, mx.array]],
        rope_offset: int,
        slot: Optional[_KVSlot],
    ) -> mx.array:
        x = x + self.attn(mxf.rms_norm(x, weight=None, eps=self.eps), mask, rope_offset, slot)
        x = x + self.mlp(mxf.rms_norm(x, weight=None, eps=self.eps))
        return x


class HrmStack(nn.Module):
    """One transformer stack — used twice as L_module and H_module."""

    def __init__(self, cfg: HrmTextConfig):
        super().__init__()
        self.layers = [HrmLayer(cfg) for _ in range(cfg.num_layers_per_stack)]
        self.eps = cfg.rms_norm_eps

    def __call__(
        self,
        x: mx.array,
        mask: Optional[Union[str, mx.array]],
        rope_offset: int,
        cache: Optional[HrmKVCache],
        slot_base: int,
    ) -> mx.array:
        for i, layer in enumerate(self.layers):
            slot = cache.slot(slot_base + i) if cache is not None else None
            x = layer(x, mask, rope_offset, slot)
        return mxf.rms_norm(x, weight=None, eps=self.eps)


# ---------------------------------------------------------------------------
# Top-level model
# ---------------------------------------------------------------------------
class HrmTextModel(nn.Module):
    def __init__(self, cfg: HrmTextConfig):
        super().__init__()
        self.cfg = cfg
        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        self.L_module = HrmStack(cfg)
        self.H_module = HrmStack(cfg)

    def make_cache(self, step: int = 256) -> HrmKVCache:
        return HrmKVCache(self.cfg, step=step)

    def __call__(
        self,
        input_ids: mx.array,
        cache: Optional[HrmKVCache] = None,
        token_type_ids: Optional[mx.array] = None,
    ) -> mx.array:
        B, S = input_ids.shape
        rope_offset = cache.offset if cache is not None else 0

        z_H = self.embed_tokens(input_ids) * self.cfg.embedding_scale
        z_L = mx.zeros_like(z_H)

        mask = self._build_mask(S, rope_offset, token_type_ids, dtype=z_H.dtype)

        num_per_stack = self.cfg.num_layers_per_stack
        L_plus_1 = self.cfg.L_cycles + 1

        for h in range(self.cfg.H_cycles):
            for l in range(self.cfg.L_cycles):
                slot_base = (h * L_plus_1 + l) * num_per_stack
                z_L = self.L_module(z_L + z_H, mask, rope_offset, cache, slot_base)
            slot_base = (h * L_plus_1 + self.cfg.L_cycles) * num_per_stack
            z_H = self.H_module(z_H + z_L, mask, rope_offset, cache, slot_base)

        return z_H

    def _build_mask(
        self,
        S: int,
        rope_offset: int,
        token_type_ids: Optional[mx.array],
        dtype,
    ) -> Optional[Union[str, mx.array]]:
        """Returns attention mask for fast.scaled_dot_product_attention.

        Fast path: prompt is all-prefix (token_type_ids=ones) on first call → full
        attention (return None). Single-token step (S==1) → return None (causal trivially).

        General path: build explicit additive mask of shape (1, 1, S, S+offset).
        """
        if S == 1:
            return None
        if token_type_ids is None or self._is_all_prefix(token_type_ids):
            # First call, fully bidirectional within prefix block.
            return None

        # General PrefixLM mask: positions i in prefix and j in prefix attend mutually.
        # Otherwise causal (j <= i).
        idx = mx.arange(S)
        causal = idx[:, None] >= idx[None, :]
        tti = token_type_ids[0] == 1
        prefix = tti[:, None] & tti[None, :]
        allow = causal | prefix
        m = mx.where(allow, mx.array(0.0, dtype=dtype), mx.array(-mx.inf, dtype=dtype))
        return m[None, None, :, :]

    @staticmethod
    def _is_all_prefix(token_type_ids: mx.array) -> bool:
        # token_type_ids is (B, S) of int. Check first batch's positions are all 1.
        # We avoid synchronization by using a cheap reduction.
        return bool(mx.all(token_type_ids == 1).item())


class HrmTextForCausalLM(nn.Module):
    def __init__(self, cfg: HrmTextConfig):
        super().__init__()
        self.cfg = cfg
        self.model = HrmTextModel(cfg)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)

    def __call__(
        self,
        input_ids: mx.array,
        cache: Optional[HrmKVCache] = None,
        token_type_ids: Optional[mx.array] = None,
    ) -> mx.array:
        h = self.model(input_ids, cache=cache, token_type_ids=token_type_ids)
        return self.lm_head(h)

    def make_cache(self, step: int = 256) -> HrmKVCache:
        return self.model.make_cache(step=step)


# ---------------------------------------------------------------------------
# Weight loading
# ---------------------------------------------------------------------------
def _load_safetensors_to_mlx(path: Path, dtype) -> dict:
    """Read a safetensors file containing bf16 tensors into MLX arrays.

    `safetensors` numpy/mlx backends can't decode bf16 directly, so we go through
    torch as a transient bridge. The bf16 → fp32 → mlx → bf16 round trip preserves
    values exactly for bf16 tensors (no precision loss).
    """
    import torch

    out = {}
    with safe_open(str(path), framework="pt") as f:
        for k in f.keys():
            t = f.get_tensor(k)
            out[k] = mx.array(t.float().numpy()).astype(dtype)
    return out


def _load_mlx_native_safetensors(path: Path) -> dict:
    """Read an MLX-native safetensors file (already-mlx-saved)."""
    return mx.load(str(path))


def _apply_weights(model: HrmTextForCausalLM, params: dict) -> None:
    # z_L_init is zeros and frozen; we instantiate at call time.
    params = {k: v for k, v in params.items() if k != "model.z_L_init"}
    model.load_weights(list(params.items()), strict=True)


def load(model_dir: Union[str, Path], dtype=mx.bfloat16) -> tuple[HrmTextForCausalLM, "Tokenizer"]:
    """Load HRM-Text-1B from a directory.

    The directory may contain either:
      - HF native safetensors (`model.safetensors` with torch bf16), or
      - MLX-converted safetensors (mlx.save_safetensors).

    Returns (model, tokenizer).
    """
    model_dir = Path(model_dir)
    cfg_dict = json.loads((model_dir / "config.json").read_text())
    cfg = HrmTextConfig.from_dict(cfg_dict)
    model = HrmTextForCausalLM(cfg)

    mlx_native = model_dir / "model_mlx.safetensors"
    hf_native = model_dir / "model.safetensors"
    if mlx_native.exists():
        params = _load_mlx_native_safetensors(mlx_native)
        # MLX-native already filtered keys; trust them.
        model.load_weights(list(params.items()), strict=True)
    elif hf_native.exists():
        params = _load_safetensors_to_mlx(hf_native, dtype)
        _apply_weights(model, params)
    else:
        raise FileNotFoundError(f"No safetensors found in {model_dir}")

    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
    return model, tokenizer


def convert_to_mlx(src_dir: Union[str, Path], dst_dir: Union[str, Path], dtype=mx.bfloat16) -> None:
    """One-shot conversion: HF safetensors → MLX-native safetensors.

    After conversion, `load(dst_dir)` reads the MLX file directly (no torch needed).
    """
    src_dir = Path(src_dir)
    dst_dir = Path(dst_dir)
    dst_dir.mkdir(parents=True, exist_ok=True)

    params = _load_safetensors_to_mlx(src_dir / "model.safetensors", dtype)
    params = {k: v for k, v in params.items() if k != "model.z_L_init"}
    mx.save_safetensors(str(dst_dir / "model_mlx.safetensors"), params)

    # Copy config, tokenizer files
    import shutil
    for name in ["config.json", "tokenizer.json", "tokenizer_config.json", "LICENSE", "README.md"]:
        src = src_dir / name
        if src.exists():
            shutil.copy(src, dst_dir / name)


# ---------------------------------------------------------------------------
# Sampling
# ---------------------------------------------------------------------------
def _sample(logits: mx.array, temperature: float, top_p: float, rng) -> int:
    """Sample one token id from (V,) logits."""
    if temperature == 0.0:
        return int(mx.argmax(logits).item())

    logits = logits.astype(mx.float32) / temperature
    if top_p < 1.0:
        sorted_idx = mx.argsort(-logits)
        sorted_logits = mx.take(logits, sorted_idx)
        sorted_probs = mx.softmax(sorted_logits, axis=-1)
        cumprobs = mx.cumsum(sorted_probs, axis=-1)
        # Keep tokens up to and including the first that crosses top_p.
        keep = cumprobs <= top_p
        # Always keep at least the top token.
        keep = mx.concatenate([mx.array([True]), keep[:-1]])
        sorted_logits = mx.where(keep, sorted_logits, mx.array(-mx.inf))
        # Restore original order.
        masked = mx.zeros_like(logits) + (-mx.inf)
        masked[sorted_idx] = sorted_logits
        logits = masked

    probs = mx.softmax(logits, axis=-1)
    # Categorical sampling
    token = int(mx.random.categorical(logits[None, :]).item())
    return token


# ---------------------------------------------------------------------------
# Condition tag handling
# ---------------------------------------------------------------------------
_CONDITION_TOKEN = {
    "direct": "<|object_ref_start|>",
    "cot": "<|object_ref_end|>",
    "noisy": "<|quad_start|>",
    "synth": "<|quad_end|>",
}


def format_prompt(text: str, condition: str = "synth,cot") -> str:
    """Wrap a prompt with HRM-Text's <|im_start|> ... <|im_end|> envelope and
    the given composite condition prefix.

    `condition` is a comma-separated list of tags, e.g. "synth,cot" for the
    reasoning-style composite or "direct" for direct-answer mode.
    """
    parts = [t.strip() for t in condition.split(",")]
    prefix = "".join(_CONDITION_TOKEN[p] for p in parts)
    return f"<|im_start|>{prefix}{text}<|im_end|>"


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------
def generate_step(
    model: HrmTextForCausalLM,
    prompt_ids: list[int],
    temperature: float = 0.0,
    top_p: float = 1.0,
    seed: int = 0,
) -> Iterator[int]:
    """Yields token ids one at a time.

    First yield is the first generated token (after prompt). Caller decides when
    to stop (EOS, max_tokens, etc.) by breaking out of the iterator.
    """
    if temperature > 0.0:
        mx.random.seed(seed)
    cfg = model.cfg
    cache = model.make_cache()

    # Prefill: pass the whole prompt with token_type_ids=ones → bidirectional prefix.
    iid = mx.array([prompt_ids], dtype=mx.int32)
    tti = mx.ones((1, len(prompt_ids)), dtype=mx.int32)
    logits = model(iid, cache=cache, token_type_ids=tti)
    last = logits[0, -1, :]
    token = _sample(last, temperature, top_p, None)
    yield token

    while True:
        iid = mx.array([[token]], dtype=mx.int32)
        # No token_type_ids on subsequent steps → causal for new tokens.
        logits = model(iid, cache=cache)
        last = logits[0, -1, :]
        token = _sample(last, temperature, top_p, None)
        yield token


def generate(
    model: HrmTextForCausalLM,
    tokenizer,
    prompt: str,
    condition: str = "synth,cot",
    max_tokens: int = 512,
    temperature: float = 0.0,
    top_p: float = 1.0,
    seed: int = 0,
    stream: bool = False,
    skip_special_tokens: bool = True,
) -> Union[str, Iterator[str]]:
    """High-level generation API.

    Args:
      condition: comma-separated condition tags. Common values:
          "synth,cot"  - reasoning / chain-of-thought (DEFAULT, good for math)
          "synth"      - synthetic / curated explanation
          "direct"     - direct answer, no CoT (good for multi-choice, extraction)
          "cot"        - chain-of-thought alone
          "noisy"      - noisy / web-crawl style
      stream: if True returns an iterator that yields decoded text chunks.

    Returns:
      If stream=False: the decoded completion (str).
      If stream=True: iterator of decoded text fragments.
    """
    formatted = format_prompt(prompt, condition)
    prompt_ids = tokenizer.encode(formatted, add_special_tokens=False)
    eos = model.cfg.eos_token_id

    def _gen_ids():
        for i, tok in enumerate(generate_step(model, prompt_ids, temperature, top_p, seed)):
            if i >= max_tokens:
                return
            if eos is not None and tok == eos:
                return
            yield tok

    if stream:
        def _stream():
            buf = []
            for tok in _gen_ids():
                buf.append(tok)
                # Decode incrementally — sometimes a token mid-multibyte produces "" which is fine.
                text = tokenizer.decode([tok], skip_special_tokens=skip_special_tokens)
                yield text
        return _stream()

    ids = list(_gen_ids())
    return tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)