Publish Modilify Mk1 MLX runtime

README.md

@@ -0,0 +1,59 @@
+---
+license: other
+license_name: modilify-open-model-license-1.0
+license_link: LICENSE
+library_name: mlx
+pipeline_tag: image-text-to-text
+base_model: modilify/Modilify-Mk1
+tags:
+  - mlx
+  - diffusion
+  - multimodal
+  - image-text-to-text
+  - mixture-of-experts
+---
+
+![LOGO](assets/01-LOGO.jpg)
+
+# Modilify Mk1 MLX
+
+Native MLX runtime for [Modilify Mk1](https://huggingface.co/modilify/Modilify-Mk1).
+
+`model_type` is `modilify_mk1`. This package does not rewrite the checkpoint as another architecture and does not install into oMLX.
+
+The published `model-*.safetensors` shards are a server-side copy of the official Mk1 weights. `modilify_mlx.load` remaps PyTorch parameter names at load time, so you do not need a second converted 49GB checkpoint on the Hub.
+
+## Layout
+
+- `modilify_mlx/` — config, latent deliberation, exclusive commit policy, trunk wrapper, generate loop
+- `convert_modilify_mk1_mlx.py` — official safetensors → MLX shards
+- `generate_modilify.py` — standalone text trial inference
+
+The heavy decoder / vision trunk is composed from mlx-vlm's DiffusionGemma layers as an internal dependency. Router top-k and the RMS-capped latent merge are Mk1's own forwards.
+
+## Convert
+
+Use the oMLX bundled interpreter only as a Python that already has `mlx`:
+
+```bash
+export PYTHONHOME=/Applications/oMLX.app/Contents/Resources/Python/cpython-3.11
+export PYTHONPATH=/Applications/oMLX.app/Contents/Resources:/Applications/oMLX.app/Contents/Resources/Python/framework-mlx-base/lib/python3.11/site-packages
+/Applications/oMLX.app/Contents/Resources/Python/cpython-3.11/bin/python3 \
+  ~/Modilify-Mk1-MLX/convert_modilify_mk1_mlx.py \
+  --source ~/Modilify-Mk1 \
+  --destination ~/Modilify-Mk1-MLX
+```
+
+## Generate
+
+Official Mk1 shards and locally remapped MLX shards both load. Parameter names are remapped in `modilify_mlx.load`.
+
+```bash
+/Applications/oMLX.app/Contents/Resources/Python/cpython-3.11/bin/python3 \
+  ~/Modilify-Mk1-MLX/generate_modilify.py \
+  --model ~/Modilify-Mk1-MLX \
+  --prompt "Explain why the sky is blue." \
+  --max-new-tokens 128
+```
+
+Default is bf16 (`--expert-bits 16`). Add `--profile` to print per-phase denoise timings. Full-graph `mx.compile` of the decoder is blocked by MoE gather/router slices.

config.json

@@ -0,0 +1,149 @@
+{
+  "architectures": [
+    "ModilifyMk1ForBlockDiffusion"
+  ],
+  "boi_token_id": 255999,
+  "bos_token_id": 2,
+  "canvas_length": 256,
+  "commit_failure_budget": 0.2,
+  "denoise_temperature": 0.8,
+  "dtype": "bfloat16",
+  "eoi_token_id": 258882,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "image_token_id": 258880,
+  "initializer_range": 0.02,
+  "jump_failure_budget": 2.0,
+  "jump_on_no_progress_after": 12,
+  "latent_dim": 1536,
+  "latent_dropout": 0.0,
+  "latent_local_attention_window": 128,
+  "latent_memory_slots": 64,
+  "latent_num_heads": 16,
+  "latent_num_layers": 4,
+  "max_ponder_steps": 64,
+  "min_trajectory_progress": 0.005,
+  "model_type": "modilify_mk1",
+  "pad_token_id": 0,
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "bos_token_id": 2,
+    "dtype": "bfloat16",
+    "eos_token_id": 1,
+    "final_logit_softcapping": 30.0,
+    "global_head_dim": 512,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 2816,
+    "initializer_range": 0.02,
+    "intermediate_size": 2112,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "model_type": "modilify_mk1_text",
+    "moe_intermediate_size": 704,
+    "num_attention_heads": 16,
+    "num_experts": 128,
+    "num_global_key_value_heads": 2,
+    "num_hidden_layers": 30,
+    "num_key_value_heads": 8,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 1000000.0,
+        "rope_type": "proportional"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 1024,
+    "tie_word_embeddings": true,
+    "top_k_experts": 8,
+    "use_bidirectional_attention": "vision",
+    "vocab_size": 262144
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.14.1",
+  "turn_end_token_id": 106,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "chunk_size_feed_forward": 0,
+    "default_output_length": 280,
+    "dtype": "bfloat16",
+    "global_head_dim": 72,
+    "head_dim": 72,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "max_position_embeddings": 131072,
+    "model_type": "gemma4_vision",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 27,
+    "num_key_value_heads": 16,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "position_embedding_size": 10240,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 100.0,
+      "rope_type": "default"
+    },
+    "standardize": true,
+    "use_clipped_linears": false
+  },
+  "repetition_penalty": 1.0
+}

convert_modilify_mk1_mlx.py

@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Convert an official Modilify Mk1 checkpoint into native MLX shards."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+
+ROOT = Path(__file__).resolve().parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from modilify_mlx.config import MODEL_TYPE, TEXT_MODEL_TYPE, ModilifyMk1Config
+from modilify_mlx.convert_utils import remap_weight, should_keep_source_key
+
+COPY_PATTERNS = (
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "chat_template.jinja",
+    "processor_config.json",
+    "generation_config.json",
+    "special_tokens_map.json",
+    "preprocessor_config.json",
+)
+
+MAX_SHARD_BYTES = 5 * 1024**3
+
+
+def _load_json(path: Path) -> dict:
+    with path.open(encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def _write_config(source: Path, destination: Path) -> None:
+    payload = _load_json(source / "config.json")
+    if payload.get("model_type") != MODEL_TYPE:
+        raise ValueError(
+            f"Source model_type must be {MODEL_TYPE!r}, got {payload.get('model_type')!r}."
+        )
+    payload.pop("auto_map", None)
+    payload["architectures"] = ["ModilifyMk1ForBlockDiffusion"]
+    text = dict(payload.get("text_config") or {})
+    text["model_type"] = TEXT_MODEL_TYPE
+    payload["text_config"] = text
+    payload["model_type"] = MODEL_TYPE
+    # Validate through the native config so a bad export fails here.
+    ModilifyMk1Config.from_dict(payload)
+    with (destination / "config.json").open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2)
+        handle.write("\n")
+
+
+def _copy_sidecar_files(source: Path, destination: Path) -> None:
+    for name in COPY_PATTERNS:
+        src = source / name
+        if src.exists():
+            shutil.copy2(src, destination / name)
+
+
+def _flush_shard(
+    destination: Path,
+    shard: dict[str, mx.array],
+    shard_index: int,
+    planned_count: int,
+    weight_map: dict[str, str],
+) -> int:
+    if not shard:
+        return shard_index
+    name = f"model-{shard_index:05d}-of-{planned_count:05d}.safetensors"
+    mx.save_safetensors(
+        str(destination / name),
+        shard,
+        metadata={"format": "mlx"},
+    )
+    for key in shard:
+        weight_map[key] = name
+    return shard_index + 1
+
+
+def convert(source: Path, destination: Path) -> None:
+    source = source.resolve()
+    destination = destination.resolve()
+    destination.mkdir(parents=True, exist_ok=True)
+    index = _load_json(source / "model.safetensors.index.json")
+    source_map: dict[str, str] = index["weight_map"]
+    shard_names = []
+    for name in source_map.values():
+        if name not in shard_names:
+            shard_names.append(name)
+
+    remapped: dict[str, mx.array] = {}
+    kept_source = 0
+    dropped_source = 0
+    for shard_name in shard_names:
+        print(f"[convert] reading {shard_name}", flush=True)
+        loaded = mx.load(str(source / shard_name))
+        for key, value in loaded.items():
+            if not should_keep_source_key(key):
+                dropped_source += 1
+                continue
+            kept_source += 1
+            for new_key, new_value in remap_weight(key, value):
+                if new_key in remapped:
+                    raise ValueError(f"Duplicate remapped key: {new_key}")
+                remapped[new_key] = new_value
+        del loaded
+
+    print(
+        f"[convert] kept {kept_source} source tensors, dropped {dropped_source}, "
+        f"wrote {len(remapped)} MLX tensors",
+        flush=True,
+    )
+    latent_keys = [key for key in remapped if key.startswith("latent_deliberation.")]
+    if len(latent_keys) < 150:
+        raise RuntimeError(
+            f"Expected the latent stack to survive conversion, found {len(latent_keys)} keys."
+        )
+
+    items = sorted(remapped.items())
+    shards: list[dict[str, mx.array]] = []
+    current: dict[str, mx.array] = {}
+    current_bytes = 0
+    for key, value in items:
+        tensor_bytes = int(value.nbytes)
+        if current and current_bytes + tensor_bytes > MAX_SHARD_BYTES:
+            shards.append(current)
+            current = {}
+            current_bytes = 0
+        current[key] = value
+        current_bytes += tensor_bytes
+    if current:
+        shards.append(current)
+
+    planned = max(1, len(shards))
+    weight_map: dict[str, str] = {}
+    total_size = 0
+    for index_i, shard in enumerate(shards, start=1):
+        print(f"[convert] writing shard {index_i}/{planned}", flush=True)
+        name = (
+            f"model-{index_i:05d}-of-{planned:05d}.safetensors"
+            if planned > 1
+            else "model.safetensors"
+        )
+        mx.save_safetensors(
+            str(destination / name),
+            shard,
+            metadata={"format": "mlx"},
+        )
+        for key, value in shard.items():
+            weight_map[key] = name
+            total_size += int(value.nbytes)
+
+    index_payload = {
+        "metadata": {"total_size": total_size},
+        "weight_map": {key: weight_map[key] for key in sorted(weight_map)},
+    }
+    with (destination / "model.safetensors.index.json").open(
+        "w", encoding="utf-8"
+    ) as handle:
+        json.dump(index_payload, handle, indent=2)
+        handle.write("\n")
+
+    _write_config(source, destination)
+    _copy_sidecar_files(source, destination)
+    print(f"[convert] done -> {destination}", flush=True)
+    print(f"[convert] tensors={len(weight_map)} bytes={total_size}", flush=True)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--source",
+        type=Path,
+        default=Path.home() / "Modilify-Mk1",
+        help="Official Mk1 Hugging Face directory",
+    )
+    parser.add_argument(
+        "--destination",
+        type=Path,
+        default=ROOT,
+        help="Native MLX output directory",
+    )
+    args = parser.parse_args()
+    convert(args.source, args.destination)
+
+
+if __name__ == "__main__":
+    main()

generate_modilify.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Standalone text trial-inference CLI for native Modilify Mk1 MLX."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+
+ROOT = Path(__file__).resolve().parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from modilify_mlx.generate import generate
+from modilify_mlx.modeling import load
+
+
+def _build_prompt_ids(model_path: Path, prompt: str, enable_thinking: bool):
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True)
+    messages = [{"role": "user", "content": prompt}]
+    token_ids = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        enable_thinking=enable_thinking,
+    )
+    if hasattr(token_ids, "input_ids"):
+        token_ids = token_ids.input_ids
+    elif isinstance(token_ids, dict):
+        token_ids = token_ids["input_ids"]
+    if hasattr(token_ids, "tolist"):
+        token_ids = token_ids.tolist()
+    if token_ids and isinstance(token_ids[0], (list, tuple)):
+        token_ids = token_ids[0]
+    return [int(token) for token in token_ids], tokenizer
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", type=Path, default=ROOT)
+    parser.add_argument("--prompt", default="Explain why the sky is blue.")
+    parser.add_argument("--max-new-tokens", type=int, default=128)
+    parser.add_argument("--temperature", type=float, default=None)
+    parser.add_argument("--enable-thinking", action="store_true")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--expert-bits",
+        type=int,
+        default=16,
+        help="Quantize MoE experts to this bit width (16 keeps bf16).",
+    )
+    parser.add_argument(
+        "--profile",
+        action="store_true",
+        help="Print per-phase denoise timings.",
+    )
+    args = parser.parse_args()
+
+    print(f"[mk1] loading {args.model}", flush=True)
+    model, config = load(args.model, expert_bits=args.expert_bits)
+    if config.model_type != "modilify_mk1":
+        raise SystemExit(f"Refusing model_type={config.model_type!r}")
+    prompt_ids, tokenizer = _build_prompt_ids(
+        args.model, args.prompt, args.enable_thinking
+    )
+    print(
+        f"[mk1] prompt_tokens={len(prompt_ids)} canvas={config.canvas_length} "
+        f"temp={args.temperature if args.temperature is not None else config.denoise_temperature}",
+        flush=True,
+    )
+    profiler = None
+    if args.profile:
+        from modilify_mlx.profile import DenoiseProfiler
+
+        profiler = DenoiseProfiler()
+    output = generate(
+        model,
+        mx.array([prompt_ids], dtype=mx.int32),
+        max_new_tokens=args.max_new_tokens,
+        temperature=args.temperature,
+        seed=args.seed,
+        profiler=profiler,
+    )
+    text = tokenizer.decode(output.generated_ids, skip_special_tokens=False)
+    visible = tokenizer.decode(output.generated_ids, skip_special_tokens=True)
+    print(
+        f"[mk1] stop={output.stop_reason} denoise={output.denoise_steps} "
+        f"committed={output.generated_length} avg_commit={output.average_commit_len:.2f} "
+        f"tpf={output.tokens_per_forward:.2f} jumps={output.jump_count}",
+        flush=True,
+    )
+    print(
+        f"[mk1] prefill={output.prefill_seconds:.3f}s "
+        f"first_denoise={output.first_denoise_seconds:.3f}s "
+        f"generate={output.generate_seconds:.3f}s "
+        f"hd/s={output.heavy_denoise_per_second:.3f} "
+        f"steady_hd/s={output.steady_heavy_denoise_per_second:.3f} "
+        f"tok/s={output.tokens_per_second:.2f}",
+        flush=True,
+    )
+    if profiler is not None:
+        print(profiler.summary(), flush=True)
+    print("--- raw ---")
+    print(text)
+    print("--- visible ---")
+    print(visible)
+
+
+if __name__ == "__main__":
+    main()

modilify_mlx/__init__.py

@@ -0,0 +1,12 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Native MLX runtime for Modilify Mk1."""
+
+from .config import ModilifyMk1Config
+from .modeling import ModilifyMk1ForBlockDiffusion, load
+
+__all__ = [
+    "ModilifyMk1Config",
+    "ModilifyMk1ForBlockDiffusion",
+    "load",
+]

modilify_mlx/commit_policy.py

@@ -0,0 +1,273 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Confidence-and-entropy commit policy for Mk1 inference.
+
+PyTorch → MLX:
+    fused_commit_confidence          same excess-entropy sigmoid-square
+    prefix_failure_commit_lengths    longest prefix with cumsum(risk) < budget
+    first_committed_token_lengths    clip after first stop token
+    select_commit_lengths            sampled prefix or greedy jump
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+import math
+
+import mlx.core as mx
+
+from .latent_deliberation import (
+    advance_trajectory_clocks,
+    should_force_trajectory_jump,
+)
+
+FUSED_EPS = 1e-6
+
+
+def fused_commit_confidence(
+    proposal_confidence: mx.array,
+    token_entropy: mx.array,
+    *,
+    vocab_size: int = 256000,
+    eps: float = FUSED_EPS,
+) -> mx.array:
+    """Fuse proposal confidence with token entropy.
+
+    p = clamp(proposal_confidence, eps, 1 - eps)
+    h2 = -p * log(p) - (1 - p) * log(1 - p)
+    excess = max(token_entropy - h2, 0)
+    fused = sigmoid(logit(p) - excess) ** 2
+    """
+
+    del vocab_size
+    p = mx.clip(proposal_confidence.astype(mx.float32), eps, 1.0 - eps)
+    entropy = mx.maximum(token_entropy.astype(mx.float32), 0.0)
+    binary_entropy = -p * mx.log(p) - (1.0 - p) * mx.log1p(-p)
+    excess = mx.maximum(entropy - binary_entropy, 0.0)
+    logit_p = mx.log(p) - mx.log1p(-p)
+    fused = mx.square(mx.sigmoid(logit_p - excess))
+    return mx.clip(fused, eps, 1.0 - eps)
+
+
+def fused_commit_failure_rate(
+    proposal_confidence: mx.array,
+    token_entropy: mx.array,
+    **kwargs: object,
+) -> mx.array:
+    """Return ``1 - fused_commit_confidence``."""
+
+    return 1.0 - fused_commit_confidence(
+        proposal_confidence, token_entropy, **kwargs
+    )
+
+
+@dataclass(frozen=True)
+class CommitPolicyDecision:
+    """One inference transition from proposal to committed prefix."""
+
+    normal_lengths: mx.array
+    commit_lengths: mx.array
+    commit_token_ids: mx.array
+    jump_rows: mx.array
+    ponder_steps: mx.array
+    stagnation_steps: mx.array
+
+
+def prefix_failure_commit_lengths(
+    failure_rate: mx.array,
+    *,
+    failure_budget: float,
+    valid_mask: mx.array | None = None,
+) -> mx.array:
+    """Return the longest prefix with ``cumsum(failure_rate) < budget``."""
+
+    if failure_rate.ndim != 2:
+        raise ValueError("Failure rate must have shape [batch, canvas].")
+    if not math.isfinite(failure_budget) or failure_budget <= 0:
+        raise ValueError("Commit failure budget must be finite and positive.")
+    if valid_mask is None:
+        valid_mask = mx.ones(failure_rate.shape, dtype=mx.bool_)
+    if valid_mask.shape != failure_rate.shape:
+        raise ValueError("Commit validity mask must match failure rate.")
+
+    risk = mx.clip(failure_rate.astype(mx.float32), 0.0, 1.0) * valid_mask.astype(
+        mx.float32
+    )
+    cumulative_risk = mx.cumsum(risk, axis=-1)
+    contiguous_valid = mx.cumprod(valid_mask.astype(mx.int32), axis=-1).astype(
+        mx.bool_
+    )
+    allowed = (cumulative_risk < float(failure_budget)) & contiguous_valid
+    return mx.sum(mx.cumprod(allowed.astype(mx.int32), axis=-1), axis=-1)
+
+
+def first_committed_token_lengths(
+    proposal: mx.array,
+    commit_lengths: mx.array,
+    token_id: int | Sequence[int],
+) -> mx.array:
+    """Clip each prefix immediately after its first stop token."""
+
+    if proposal.ndim != 2 or commit_lengths.shape != proposal.shape[:1]:
+        raise ValueError("Proposal and commit lengths must share a batch dimension.")
+    canvas = proposal.shape[1]
+    positions = mx.arange(canvas)[None, :]
+    committed = positions < commit_lengths[:, None]
+    if isinstance(token_id, int):
+        stop_token_ids = (int(token_id),)
+    else:
+        stop_token_ids = tuple(dict.fromkeys(int(value) for value in token_id))
+    if not stop_token_ids:
+        raise ValueError("At least one stop token ID is required.")
+    matches = proposal == stop_token_ids[0]
+    for value in stop_token_ids[1:]:
+        matches = matches | (proposal == value)
+    matches = matches & committed
+    sentinel = mx.full(positions.shape, canvas, dtype=positions.dtype)
+    first = mx.min(mx.where(matches, positions, sentinel), axis=-1)
+    clipped = mx.where(first < canvas, first + 1, commit_lengths)
+    return mx.minimum(clipped, commit_lengths)
+
+
+def bounded_prefix_failure_commit_lengths(
+    committed_token_ids: mx.array,
+    failure_rate: mx.array,
+    *,
+    failure_budget: float,
+    remaining_lengths: mx.array,
+    stop_token_id: int | Sequence[int],
+    valid_mask: mx.array | None = None,
+) -> mx.array:
+    """Apply remaining-length and stop-token bounds to the prefix policy."""
+
+    if committed_token_ids.shape != failure_rate.shape:
+        raise ValueError(
+            "Committed token IDs and failure rate must share [batch, canvas]."
+        )
+    if remaining_lengths.shape != committed_token_ids.shape[:1]:
+        raise ValueError("Remaining lengths must have shape [batch].")
+    commit_lengths = prefix_failure_commit_lengths(
+        failure_rate,
+        failure_budget=failure_budget,
+        valid_mask=valid_mask,
+    )
+    commit_lengths = mx.minimum(commit_lengths, mx.maximum(remaining_lengths, 0))
+    return first_committed_token_lengths(
+        committed_token_ids,
+        commit_lengths,
+        stop_token_id,
+    )
+
+
+def select_commit_lengths(
+    sampled_token_ids: mx.array,
+    normal_failure_rate: mx.array,
+    previous_failure_rate: mx.array,
+    greedy_token_ids: mx.array,
+    jump_failure_rate: mx.array,
+    *,
+    ponder_steps: mx.array,
+    stagnation_steps: mx.array,
+    active_rows: mx.array,
+    remaining_lengths: mx.array,
+    failure_budget: float,
+    jump_failure_budget: float,
+    stop_token_id: int | Sequence[int],
+    max_ponder_steps: int,
+    stagnation_threshold: int,
+    min_progress: float,
+    valid_mask: mx.array | None = None,
+) -> CommitPolicyDecision:
+    """Select sampled commits or a greedy jump after stagnation."""
+
+    if not (
+        sampled_token_ids.shape
+        == normal_failure_rate.shape
+        == previous_failure_rate.shape
+        == greedy_token_ids.shape
+        == jump_failure_rate.shape
+    ):
+        raise ValueError("Sampled and greedy statistics must share [batch, canvas].")
+
+    normal = bounded_prefix_failure_commit_lengths(
+        sampled_token_ids,
+        normal_failure_rate,
+        failure_budget=failure_budget,
+        remaining_lengths=remaining_lengths,
+        stop_token_id=stop_token_id,
+        valid_mask=valid_mask,
+    )
+    canvas_length = normal_failure_rate.shape[1]
+    previous_prefix_length = prefix_failure_commit_lengths(
+        previous_failure_rate,
+        failure_budget=failure_budget,
+        valid_mask=valid_mask,
+    )
+    frontier_length = mx.maximum(previous_prefix_length, normal) + 1
+    if valid_mask is not None:
+        valid_lengths = mx.sum(valid_mask.astype(mx.int32), axis=-1)
+    else:
+        valid_lengths = mx.full(frontier_length.shape, canvas_length)
+    frontier_length = mx.minimum(frontier_length, valid_lengths)
+    positions = mx.arange(canvas_length)[None, :]
+    progress_mask = positions < frontier_length[:, None]
+    if valid_mask is not None:
+        progress_mask = progress_mask & valid_mask
+    progress_mask = progress_mask & active_rows[:, None]
+    signed_improvement = previous_failure_rate.astype(mx.float32) - (
+        normal_failure_rate.astype(mx.float32)
+    )
+    weights = progress_mask.astype(mx.float32)
+    progress = mx.sum(signed_improvement * weights, axis=-1) / mx.maximum(
+        mx.sum(weights, axis=-1), 1.0
+    )
+    next_ponder, next_stagnation = advance_trajectory_clocks(
+        ponder_steps,
+        stagnation_steps,
+        commit_lengths=normal,
+        active_rows=active_rows,
+        progress_scores=progress,
+        min_progress=min_progress,
+    )
+    jump_rows = (
+        (normal == 0)
+        & active_rows
+        & should_force_trajectory_jump(
+            next_ponder,
+            next_stagnation,
+            max_ponder_steps=max_ponder_steps,
+            stagnation_threshold=stagnation_threshold,
+        )
+    )
+    jump_commit = bounded_prefix_failure_commit_lengths(
+        greedy_token_ids,
+        jump_failure_rate,
+        failure_budget=jump_failure_budget,
+        remaining_lengths=remaining_lengths,
+        stop_token_id=stop_token_id,
+        valid_mask=valid_mask,
+    )
+    committed = mx.where(jump_rows, jump_commit, normal)
+    commit_token_ids = mx.where(
+        jump_rows[:, None],
+        greedy_token_ids,
+        sampled_token_ids,
+    )
+    committed = first_committed_token_lengths(
+        commit_token_ids,
+        committed,
+        stop_token_id,
+    )
+    committed = mx.where(active_rows, committed, 0)
+    jump_rows = jump_rows & (committed > 0)
+    next_ponder = mx.where(committed > 0, 0, next_ponder).astype(mx.int32)
+    next_stagnation = mx.where(committed > 0, 0, next_stagnation).astype(mx.int32)
+    return CommitPolicyDecision(
+        normal_lengths=normal,
+        commit_lengths=committed,
+        commit_token_ids=commit_token_ids,
+        jump_rows=jump_rows,
+        ponder_steps=next_ponder,
+        stagnation_steps=next_stagnation,
+    )

modilify_mlx/config.py

@@ -0,0 +1,165 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Serializable configuration for the native Modilify Mk1 MLX runtime."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import json
+from pathlib import Path
+from typing import Any
+
+from mlx_vlm.models.diffusion_gemma.config import ModelConfig, TextConfig
+from mlx_vlm.models.gemma4.config import VisionConfig
+from mlx_vlm.models.qwen3_vl.config import _config_kwargs, _maybe_deserialize_config
+
+MODEL_TYPE = "modilify_mk1"
+TEXT_MODEL_TYPE = "modilify_mk1_text"
+
+
+def _require_model_type(payload: dict[str, Any], expected: str) -> None:
+    actual = payload.get("model_type")
+    if actual != expected:
+        raise ValueError(
+            f"Expected model_type {expected!r}, got {actual!r}. "
+            "This runtime only loads native Modilify Mk1 checkpoints."
+        )
+
+
+@dataclass
+class ModilifyMk1Config:
+    """Mk1 inference configuration. ``model_type`` is always ``modilify_mk1``."""
+
+    text_config: TextConfig
+    vision_config: VisionConfig | None = None
+    model_type: str = MODEL_TYPE
+    boi_token_id: int = 255999
+    eoi_token_id: int = 258882
+    image_token_id: int = 258880
+    video_token_id: int | None = None
+    pad_token_id: int = 0
+    bos_token_id: int = 2
+    eos_token_id: list[int] | int = field(default_factory=lambda: [1, 106])
+    initializer_range: float = 0.02
+    canvas_length: int = 256
+    denoise_temperature: float = 0.8
+    commit_failure_budget: float = 0.2
+    jump_failure_budget: float = 2.0
+    jump_on_no_progress_after: int = 12
+    max_ponder_steps: int = 64
+    min_trajectory_progress: float = 0.005
+    latent_dim: int = 1536
+    latent_dropout: float = 0.0
+    latent_local_attention_window: int = 128
+    latent_memory_slots: int = 64
+    latent_num_heads: int = 16
+    latent_num_layers: int = 4
+    turn_end_token_id: int = 106
+    repetition_penalty: float = 1.0
+    dtype: str | None = "bfloat16"
+    generation_config: dict[str, Any] | None = None
+
+    def __post_init__(self) -> None:
+        if self.model_type != MODEL_TYPE:
+            raise ValueError(
+                f"ModilifyMk1Config.model_type must be {MODEL_TYPE!r}, "
+                f"got {self.model_type!r}."
+            )
+        if isinstance(self.text_config, dict):
+            text_payload = dict(self.text_config)
+            text_payload.setdefault("model_type", TEXT_MODEL_TYPE)
+            self.text_config = TextConfig(
+                **_config_kwargs(TextConfig, text_payload)
+            )
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> "ModilifyMk1Config":
+        payload = dict(payload)
+        _require_model_type(payload, MODEL_TYPE)
+        text_payload = dict(payload.get("text_config") or {})
+        text_payload.setdefault("model_type", TEXT_MODEL_TYPE)
+        payload["text_config"] = TextConfig(
+            **_config_kwargs(TextConfig, text_payload)
+        )
+        if payload.get("vision_config") is not None:
+            payload["vision_config"] = _maybe_deserialize_config(
+                VisionConfig, payload.get("vision_config")
+            )
+        return cls(**_config_kwargs(cls, payload))
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> "ModilifyMk1Config":
+        with Path(path).open(encoding="utf-8") as handle:
+            return cls.from_dict(json.load(handle))
+
+    def to_dict(self) -> dict[str, Any]:
+        def _as_dict(value: Any) -> Any:
+            if hasattr(value, "to_dict"):
+                return value.to_dict()
+            if hasattr(value, "__dict__") and not isinstance(
+                value, (str, int, float, bool, list, dict, type(None))
+            ):
+                return {
+                    key: _as_dict(item)
+                    for key, item in vars(value).items()
+                    if not key.startswith("_")
+                }
+            return value
+
+        return {
+            "model_type": MODEL_TYPE,
+            "text_config": _as_dict(self.text_config),
+            "vision_config": None
+            if self.vision_config is None
+            else _as_dict(self.vision_config),
+            "boi_token_id": self.boi_token_id,
+            "eoi_token_id": self.eoi_token_id,
+            "image_token_id": self.image_token_id,
+            "video_token_id": self.video_token_id,
+            "pad_token_id": self.pad_token_id,
+            "bos_token_id": self.bos_token_id,
+            "eos_token_id": self.eos_token_id,
+            "initializer_range": self.initializer_range,
+            "canvas_length": self.canvas_length,
+            "denoise_temperature": self.denoise_temperature,
+            "commit_failure_budget": self.commit_failure_budget,
+            "jump_failure_budget": self.jump_failure_budget,
+            "jump_on_no_progress_after": self.jump_on_no_progress_after,
+            "max_ponder_steps": self.max_ponder_steps,
+            "min_trajectory_progress": self.min_trajectory_progress,
+            "latent_dim": self.latent_dim,
+            "latent_dropout": self.latent_dropout,
+            "latent_local_attention_window": self.latent_local_attention_window,
+            "latent_memory_slots": self.latent_memory_slots,
+            "latent_num_heads": self.latent_num_heads,
+            "latent_num_layers": self.latent_num_layers,
+            "turn_end_token_id": self.turn_end_token_id,
+            "repetition_penalty": self.repetition_penalty,
+            "dtype": self.dtype,
+        }
+
+    def trunk_model_config(self) -> ModelConfig:
+        """Build the borrowed trunk config object used only inside this package."""
+
+        return ModelConfig(
+            text_config=self.text_config,
+            vision_config=self.vision_config,
+            model_type=self.model_type,
+            boi_token_id=self.boi_token_id,
+            eoi_token_id=self.eoi_token_id,
+            image_token_id=self.image_token_id,
+            video_token_id=self.video_token_id,
+            initializer_range=self.initializer_range,
+            canvas_length=self.canvas_length,
+            eos_token_id=self.eos_token_id,
+            generation_config=self.generation_config,
+            dtype=self.dtype,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.text_config.vocab_size)
+
+    @property
+    def hidden_size(self) -> int:
+        return int(self.text_config.hidden_size)

modilify_mlx/convert_utils.py

@@ -0,0 +1,124 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Weight remapping from the official Mk1 safetensors export to MLX."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import mlx.core as mx
+
+# PyTorch → MLX
+#   *.experts.down_proj              *.experts.down_proj.weight
+#   *.experts.gate_up_proj           *.experts.gate_up_proj.weight
+#   *.in_proj_weight [3D, D]         query/key/value_proj.weight
+#   *.in_proj_bias [3D]              query/key/value_proj.bias
+#   token_ff.0 / token_ff.2          token_ff.layers.0 / token_ff.layers.2
+#   memory_ff.0 / memory_ff.2        memory_ff.layers.0 / memory_ff.layers.2
+
+_ATTENTION_MODULES = (
+    "local_attention",
+    "token_memory_attention",
+    "memory_token_attention",
+)
+
+_SKIP_SUBSTRINGS = (
+    "rotary_emb",
+    "lm_head.weight",
+)
+
+_CLIP_MARKERS = ("input_max", "input_min", "output_max", "output_min")
+
+
+def should_keep_source_key(key: str) -> bool:
+    if any(marker in key for marker in _SKIP_SUBSTRINGS):
+        return False
+    if key.startswith("model.encoder.language_model.") and not key.endswith(
+        ".layer_scalar"
+    ):
+        return False
+    if key.startswith("model.encoder.vision_tower.") or key.startswith(
+        "model.encoder.embed_vision."
+    ):
+        if any(marker in key for marker in _CLIP_MARKERS):
+            return False
+    return True
+
+
+def _split_qkv(prefix: str, value: mx.array) -> list[tuple[str, mx.array]]:
+    if value.ndim == 1:
+        width = value.shape[0]
+        if width % 3:
+            raise ValueError(f"Cannot split QKV bias for {prefix}: shape {value.shape}")
+        head = width // 3
+        pieces = (value[:head], value[head : 2 * head], value[2 * head :])
+        names = ("query_proj.bias", "key_proj.bias", "value_proj.bias")
+    elif value.ndim == 2:
+        width = value.shape[0]
+        if width % 3:
+            raise ValueError(
+                f"Cannot split QKV weight for {prefix}: shape {value.shape}"
+            )
+        head = width // 3
+        pieces = (value[:head], value[head : 2 * head], value[2 * head :])
+        names = ("query_proj.weight", "key_proj.weight", "value_proj.weight")
+        if pieces[0].shape[0] != pieces[0].shape[1]:
+            raise ValueError(
+                f"Split QKV weight for {prefix} is not square: {pieces[0].shape}"
+            )
+    else:
+        raise ValueError(f"Unexpected QKV tensor rank for {prefix}: {value.shape}")
+    return [(f"{prefix}.{name}", piece) for name, piece in zip(names, pieces)]
+
+
+def remap_weight(key: str, value: mx.array) -> list[tuple[str, mx.array]]:
+    """Map one official Mk1 tensor onto one or more MLX parameter names."""
+
+    if not should_keep_source_key(key):
+        return []
+
+    if key.endswith(".experts.down_proj"):
+        return [(key + ".weight", value)]
+    if key.endswith(".experts.gate_up_proj"):
+        return [(key + ".weight", value)]
+
+    for module in _ATTENTION_MODULES:
+        in_proj_weight = f".{module}.in_proj_weight"
+        in_proj_bias = f".{module}.in_proj_bias"
+        if key.endswith(in_proj_weight):
+            prefix = key[: -len(".in_proj_weight")]
+            return _split_qkv(prefix, value)
+        if key.endswith(in_proj_bias):
+            prefix = key[: -len(".in_proj_bias")]
+            return _split_qkv(prefix, value)
+
+    if ".token_ff.0." in key or key.endswith(".token_ff.0.weight") or key.endswith(
+        ".token_ff.0.bias"
+    ):
+        return [(key.replace(".token_ff.0.", ".token_ff.layers.0."), value)]
+    if ".token_ff.2." in key or key.endswith(".token_ff.2.weight") or key.endswith(
+        ".token_ff.2.bias"
+    ):
+        return [(key.replace(".token_ff.2.", ".token_ff.layers.2."), value)]
+    if ".memory_ff.0." in key or key.endswith(".memory_ff.0.weight") or key.endswith(
+        ".memory_ff.0.bias"
+    ):
+        return [(key.replace(".memory_ff.0.", ".memory_ff.layers.0."), value)]
+    if ".memory_ff.2." in key or key.endswith(".memory_ff.2.weight") or key.endswith(
+        ".memory_ff.2.bias"
+    ):
+        return [(key.replace(".memory_ff.2.", ".memory_ff.layers.2."), value)]
+
+    return [(key, value)]
+
+
+def remap_state_dict(
+    source: Iterable[tuple[str, mx.array]],
+) -> dict[str, mx.array]:
+    remapped: dict[str, mx.array] = {}
+    for key, value in source:
+        for new_key, new_value in remap_weight(key, value):
+            if new_key in remapped:
+                raise ValueError(f"Duplicate remapped key: {new_key}")
+            remapped[new_key] = new_value
+    return remapped

modilify_mlx/fast_decode.py

@@ -0,0 +1,298 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Compile-friendly Mk1 decoder: static prefix KV + fixed-shape sliding gather."""
+
+from __future__ import annotations
+
+from typing import Any
+import time
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_vlm.models.base import scaled_dot_product_attention
+from mlx_vlm.models.diffusion_gemma.language import geglu
+
+
+def make_static_cache(encoder, max_size: int):
+    """Fixed-capacity prefix cache so decoder_state has a static key length."""
+
+    return encoder.make_cache(max_size=max(int(max_size), 1))
+
+
+def _cache_capacity(cache) -> int:
+    first = cache[0]
+    keys = getattr(first, "keys", None)
+    if keys is None:
+        return int(getattr(first, "max_size", 0) or 0)
+    return int(keys.shape[2])
+
+
+def _prefix_length(cache) -> int:
+    offset = getattr(cache[0], "offset", 0)
+    if isinstance(offset, mx.array):
+        return int(mx.max(offset).item())
+    return int(offset)
+
+
+def build_decoder_masks(
+    *,
+    prefix_len: int,
+    canvas_length: int,
+    cache_capacity: int,
+    sliding_window: int,
+    batch_size: int = 1,
+) -> tuple[mx.array, mx.array]:
+    """Boolean SDPA masks with static shapes for full and sliding layers."""
+
+    cache_valid = mx.arange(cache_capacity) < int(prefix_len)
+    canvas_valid = mx.ones((canvas_length,), dtype=mx.bool_)
+    full_row = mx.concatenate([cache_valid, canvas_valid], axis=0)
+    full = mx.broadcast_to(
+        full_row.reshape(1, 1, 1, -1),
+        (batch_size, 1, canvas_length, cache_capacity + canvas_length),
+    )
+
+    window = max(int(sliding_window) - 1, 1)
+    slide_k = min(window, cache_capacity)
+    if slide_k < cache_capacity:
+        raw_idx = mx.arange(slide_k) + (int(prefix_len) - slide_k)
+        slide_valid = (raw_idx >= 0) & (raw_idx < int(prefix_len))
+    else:
+        slide_valid = mx.arange(cache_capacity) < int(prefix_len)
+    slide_row = mx.concatenate([slide_valid, canvas_valid], axis=0)
+    slide = mx.broadcast_to(
+        slide_row.reshape(1, 1, 1, -1),
+        (batch_size, 1, canvas_length, slide_k + canvas_length),
+    )
+    return full, slide
+
+
+def _gather_window(tensor: mx.array, offset: mx.array, window: int) -> mx.array:
+    max_size = tensor.shape[2]
+    idx = mx.clip(mx.arange(window) + offset - window, 0, max_size - 1)
+    gather = idx.reshape(1, 1, window, 1)
+    return mx.take_along_axis(tensor, gather, axis=2)
+
+
+def _decoder_attention(
+    attn: nn.Module,
+    x: mx.array,
+    mask: mx.array | None,
+    cache,
+    offset: mx.array,
+) -> mx.array:
+    batch, length, _ = x.shape
+    queries = attn.q_proj(x).reshape(batch, length, attn.n_heads, attn.head_dim)
+    queries = attn.q_norm(queries).transpose(0, 2, 1, 3)
+    queries = attn.rope(queries, offset=offset)
+
+    keys = attn.k_proj(x).reshape(batch, length, attn.n_kv_heads, attn.head_dim)
+    values = (
+        attn.v_proj(x).reshape(batch, length, attn.n_kv_heads, attn.head_dim)
+        if attn.v_proj is not None
+        else keys
+    )
+    keys = attn.k_norm(keys).transpose(0, 2, 1, 3)
+    keys = attn.rope(keys, offset=offset)
+    values = attn.v_norm(values).transpose(0, 2, 1, 3)
+
+    encoder_keys, encoder_values = cache.decoder_state
+    if attn.is_sliding:
+        window = max(int(attn.config.sliding_window) - 1, 1)
+        slide_k = min(window, int(encoder_keys.shape[2]))
+        if slide_k < int(encoder_keys.shape[2]):
+            encoder_keys = _gather_window(encoder_keys, offset, slide_k)
+            encoder_values = _gather_window(encoder_values, offset, slide_k)
+    keys = mx.concatenate([encoder_keys, keys], axis=2)
+    values = mx.concatenate([encoder_values, values], axis=2)
+    output = scaled_dot_product_attention(
+        queries, keys, values, cache=None, scale=attn.scale, mask=mask
+    )
+    output = output.transpose(0, 2, 1, 3).reshape(batch, length, -1)
+    return attn.o_proj(output)
+
+
+def _experts_unsorted(experts: nn.Module, x: mx.array, top_k_indices, top_k_weights):
+    """Expert FFN without argsort gather, so the decoder graph can compile."""
+
+    x = mx.expand_dims(x, (-2, -3))
+    gate_up = experts.gate_up_proj(x, top_k_indices, sorted_indices=False)
+    gate = gate_up[..., : experts.hidden_dims]
+    up = gate_up[..., experts.hidden_dims :]
+    y = experts.down_proj(geglu(gate, up), top_k_indices, sorted_indices=False)
+    y = y.squeeze(-2)
+    return (y * top_k_weights[..., None]).sum(axis=-2)
+
+
+def _decoder_layer(
+    layer: nn.Module,
+    x: mx.array,
+    mask: mx.array | None,
+    cache,
+    offset: mx.array,
+) -> mx.array:
+    residual = x
+    hidden = layer.input_layernorm(x)
+    hidden = _decoder_attention(layer.self_attn, hidden, mask, cache, offset)
+    hidden = layer.post_attention_layernorm(hidden)
+    hidden = residual + hidden
+
+    residual = hidden
+    shared = layer.pre_feedforward_layernorm(hidden)
+    shared = layer.mlp(shared)
+    shared = layer.post_feedforward_layernorm_1(shared)
+
+    flat = residual.reshape(-1, residual.shape[-1])
+    top_k_indices, top_k_weights = layer.router(flat)
+    routed = layer.pre_feedforward_layernorm_2(flat)
+    routed = layer.experts(routed, top_k_indices, top_k_weights)
+    routed = routed.reshape(residual.shape)
+    routed = layer.post_feedforward_layernorm_2(routed)
+
+    hidden = layer.post_feedforward_layernorm(shared + routed)
+    return residual + hidden
+
+
+def _attn_residual(
+    layer: nn.Module,
+    x: mx.array,
+    mask: mx.array | None,
+    cache,
+    offset: mx.array,
+) -> mx.array:
+    residual = x
+    hidden = layer.input_layernorm(x)
+    hidden = _decoder_attention(layer.self_attn, hidden, mask, cache, offset)
+    hidden = layer.post_attention_layernorm(hidden)
+    return residual + hidden
+
+
+def _ffn_residual(layer: nn.Module, hidden: mx.array) -> mx.array:
+    residual = hidden
+    shared = layer.pre_feedforward_layernorm(hidden)
+    shared = layer.mlp(shared)
+    shared = layer.post_feedforward_layernorm_1(shared)
+
+    flat = residual.reshape(-1, residual.shape[-1])
+    top_k_indices, top_k_weights = layer.router(flat)
+    routed = layer.pre_feedforward_layernorm_2(flat)
+    routed = layer.experts(routed, top_k_indices, top_k_weights)
+    routed = routed.reshape(residual.shape)
+    routed = layer.post_feedforward_layernorm_2(routed)
+
+    hidden = layer.post_feedforward_layernorm(shared + routed)
+    return (residual + hidden) * layer.layer_scalar
+
+
+def make_compiled_attn_layers(decoder: nn.Module, cache) -> list:
+    """Compile attention residuals only. Expert FFNs stay eager."""
+
+    compiled = []
+    for layer, layer_cache in zip(decoder.layers, cache):
+        def _fn(x, offset, mask, _layer=layer, _cache=layer_cache):
+            return _attn_residual(_layer, x, mask, _cache, offset)
+
+        compiled.append(mx.compile(_fn, shapeless=True))
+    return compiled
+
+
+def decoder_hidden_states(
+    decoder: nn.Module,
+    canvas_ids: mx.array,
+    latent_context: mx.array,
+    cache,
+    offset: mx.array,
+    full_mask: mx.array,
+    slide_mask: mx.array,
+    compiled_attn_layers=None,
+    profiler=None,
+) -> mx.array:
+    hidden = decoder._embed_canvas(
+        canvas_ids,
+        self_conditioning_embeddings=latent_context,
+    )
+    if profiler is None:
+        for layer, layer_cache in zip(decoder.layers, cache):
+            mask = slide_mask if layer.layer_type == "sliding_attention" else full_mask
+            hidden = _decoder_layer(layer, hidden, mask, layer_cache, offset)
+            hidden = hidden * layer.layer_scalar
+        return decoder.norm(hidden)
+
+    attn_acc = 0.0
+    moe_acc = 0.0
+    for layer, layer_cache in zip(decoder.layers, cache):
+        mask = slide_mask if layer.layer_type == "sliding_attention" else full_mask
+        mx.eval(hidden)
+        started = time.perf_counter()
+        hidden = _attn_residual(layer, hidden, mask, layer_cache, offset)
+        mx.eval(hidden)
+        attn_acc += time.perf_counter() - started
+        started = time.perf_counter()
+        hidden = _ffn_residual(layer, hidden)
+        mx.eval(hidden)
+        moe_acc += time.perf_counter() - started
+    profiler.add("attn", attn_acc)
+    profiler.add("moe", moe_acc)
+    return decoder.norm(hidden)
+
+
+def make_compiled_decoder(
+    decoder: nn.Module,
+    cache,
+    softcap: float,
+):
+    """Compile canvas + latent-context → (logits, hidden)."""
+
+    def _step(
+        canvas_ids: mx.array,
+        latent_context: mx.array,
+        offset: mx.array,
+        full_mask: mx.array,
+        slide_mask: mx.array,
+    ):
+        hidden = decoder_hidden_states(
+            decoder,
+            canvas_ids,
+            latent_context,
+            cache,
+            offset,
+            full_mask,
+            slide_mask,
+        )
+        logits = decoder.embed_tokens.as_linear(hidden)
+        logits = mx.tanh(logits.astype(mx.float32) / softcap) * softcap
+        return logits, hidden
+
+    return mx.compile(_step, shapeless=True)
+
+
+def quantize_experts(model: nn.Module, *, bits: int = 8, group_size: int = 64) -> None:
+    """Quantize MoE expert projections only. Attention / embeddings stay bf16."""
+
+    if bits >= 16:
+        return
+
+    def predicate(path: str, module: nn.Module):
+        if "experts" not in path:
+            return False
+        if not hasattr(module, "to_quantized"):
+            return False
+        if not (path.endswith("gate_up_proj") or path.endswith("down_proj")):
+            return False
+        return {"group_size": group_size, "bits": int(bits), "mode": "affine"}
+
+    print(f"[mk1] quantizing experts to {bits}-bit", flush=True)
+    before = sum(arr.nbytes for _, arr in model.parameters().items()) if False else None
+    del before
+    quantized = {"count": 0}
+
+    def counting_predicate(path: str, module: nn.Module):
+        result = predicate(path, module)
+        if result:
+            quantized["count"] += 1
+        return result
+
+    nn.quantize(model, class_predicate=counting_predicate)
+    print(f"[mk1] quantized {quantized['count']} expert projections", flush=True)
+    mx.eval(model.parameters())

modilify_mlx/generate.py

@@ -0,0 +1,388 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Rolling canvas generation for native Modilify Mk1."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+import math
+import time
+from typing import Any
+
+import mlx.core as mx
+
+from mlx_vlm.generate.common import wired_limit
+
+from .commit_policy import fused_commit_failure_rate, select_commit_lengths
+from .latent_deliberation import LatentDeliberationState
+from .modeling import ModilifyMk1ForBlockDiffusion
+
+
+@dataclass
+class ModilifyMk1GenerationOutput:
+    sequences: list[int]
+    generated_ids: list[int]
+    generated_length: int
+    denoise_steps: int
+    jump_count: int
+    average_commit_len: float
+    stop_reason: str
+    tokens_per_forward: float
+    prefill_seconds: float = 0.0
+    generate_seconds: float = 0.0
+    first_denoise_seconds: float = 0.0
+    heavy_denoise_per_second: float = 0.0
+    steady_heavy_denoise_per_second: float = 0.0
+    tokens_per_second: float = 0.0
+
+
+@dataclass
+class _RollingState:
+    canvas: mx.array
+    confidence: mx.array
+    entropy: mx.array
+    age: mx.array
+    latent_state: LatentDeliberationState
+    history_hidden_state: mx.array | None
+
+
+def _flatten_token_ids(*values: object) -> tuple[int, ...]:
+    token_ids: list[int] = []
+    for value in values:
+        if value is None:
+            continue
+        if isinstance(value, int):
+            token_ids.append(int(value))
+        elif isinstance(value, (list, tuple, set)):
+            token_ids.extend(int(token_id) for token_id in value if token_id is not None)
+    return tuple(dict.fromkeys(token_ids))
+
+
+def _shift_prefix(tensor: mx.array, committed: int, fill_value: float | int) -> mx.array:
+    if committed <= 0:
+        return tensor
+    canvas = tensor.shape[1]
+    if committed >= canvas:
+        return mx.full(tensor.shape, fill_value, dtype=tensor.dtype)
+    kept = tensor[:, committed:]
+    fill_shape = (tensor.shape[0], committed, *tensor.shape[2:])
+    fill = mx.full(fill_shape, fill_value, dtype=tensor.dtype)
+    return mx.concatenate([kept, fill], axis=1)
+
+
+def _shift_state(
+    state: _RollingState,
+    committed: int,
+    *,
+    vocab_size: int,
+    canvas_length: int,
+    unknown_entropy: float,
+) -> _RollingState:
+    if committed <= 0:
+        return state
+    tail = mx.random.randint(0, vocab_size, (state.canvas.shape[0], committed))
+    canvas = mx.concatenate([state.canvas[:, committed:], tail], axis=1)
+    if canvas.shape[1] != canvas_length:
+        raise RuntimeError("Canvas shift produced an unexpected length.")
+    latent = state.latent_state
+    shifted_latent = LatentDeliberationState(
+        token_latents=_shift_prefix(latent.token_latents, committed, 0),
+        memory_slots=latent.memory_slots,
+        confidence=_shift_prefix(latent.confidence, committed, 0),
+        entropy=_shift_prefix(latent.entropy, committed, unknown_entropy),
+        age=_shift_prefix(latent.age, committed, 0),
+        token_changed=_shift_prefix(latent.token_changed, committed, 0),
+        confidence_delta=_shift_prefix(latent.confidence_delta, committed, 0),
+        entropy_delta=_shift_prefix(latent.entropy_delta, committed, 0),
+        ponder_steps=mx.zeros_like(latent.ponder_steps),
+        stagnation_steps=mx.zeros_like(latent.stagnation_steps),
+    )
+    history = state.history_hidden_state
+    if history is not None:
+        history = _shift_prefix(history, committed, 0)
+    return _RollingState(
+        canvas=canvas,
+        confidence=_shift_prefix(state.confidence, committed, 0),
+        entropy=_shift_prefix(state.entropy, committed, unknown_entropy),
+        age=_shift_prefix(state.age, committed, 0),
+        latent_state=shifted_latent,
+        history_hidden_state=history,
+    )
+
+
+def generate(
+    model: ModilifyMk1ForBlockDiffusion,
+    input_ids: mx.array,
+    *,
+    max_new_tokens: int | None = None,
+    temperature: float | None = None,
+    attention_mask: mx.array | None = None,
+    pixel_values: mx.array | None = None,
+    mm_token_type_ids: mx.array | None = None,
+    max_denoising_steps: int | None = None,
+    seed: int | None = None,
+    profiler=None,
+) -> ModilifyMk1GenerationOutput:
+    """Generate one response with rolling block diffusion.
+
+    Batch size 1 only. Semantics follow ``generation_modilify_mk1.py``.
+    """
+
+    if input_ids.ndim != 2 or input_ids.shape[0] != 1:
+        raise ValueError("Native Mk1 generation currently requires shape [1, sequence].")
+    if seed is not None:
+        mx.random.seed(int(seed))
+
+    config = model.config
+    canvas_length = int(config.canvas_length)
+    vocab_size = int(config.vocab_size)
+    dtype = model.model.decoder.embed_tokens.weight.dtype
+    max_new = int(max_new_tokens if max_new_tokens is not None else 256)
+    if max_new <= 0:
+        raise ValueError("`max_new_tokens` must be positive.")
+    denoise_temperature = (
+        float(config.denoise_temperature) if temperature is None else float(temperature)
+    )
+    max_iterations = max(1, max_new * int(config.max_ponder_steps))
+    unknown_entropy = math.log(vocab_size)
+
+    if attention_mask is None:
+        attention_mask = mx.ones(input_ids.shape, dtype=mx.bool_)
+    prefill_started = time.perf_counter()
+    cache = model.make_cache(max_size=int(input_ids.shape[1]) + max_new)
+    cache = model.prefill(
+        input_ids,
+        attention_mask=attention_mask,
+        cache=cache,
+        pixel_values=pixel_values,
+        mm_token_type_ids=mm_token_type_ids,
+    )
+    mx.eval([item for block in cache for item in getattr(block, "state", ())])
+    prefill_seconds = time.perf_counter() - prefill_started
+    compiled_decoder_step = None
+
+    latent = LatentDeliberationState.empty(
+        batch_size=1,
+        canvas_length=canvas_length,
+        latent_dim=config.latent_dim,
+        memory_slots=config.latent_memory_slots,
+        dtype=dtype,
+    )
+    state = _RollingState(
+        canvas=mx.random.randint(0, vocab_size, (1, canvas_length)),
+        confidence=mx.zeros((1, canvas_length), dtype=mx.float32),
+        entropy=mx.full((1, canvas_length), unknown_entropy, dtype=mx.float32),
+        age=mx.zeros((1, canvas_length), dtype=mx.int32),
+        latent_state=latent,
+        history_hidden_state=None,
+    )
+    stop_token_ids = _flatten_token_ids(config.turn_end_token_id, config.eos_token_id)
+    turn_end = int(config.turn_end_token_id)
+    generated: list[int] = []
+    denoise_steps = 0
+    jumps = 0
+    shifts = 0
+    stop_reason = "episode_watchdog"
+    prompt_ids = [int(token) for token in input_ids[0].tolist()]
+    generate_started = time.perf_counter()
+    first_denoise_seconds = 0.0
+    always_active = mx.array([True])
+
+    with wired_limit(model, None):
+        while len(generated) < max_new:
+                step_started = time.perf_counter()
+                remaining = mx.array([max_new - len(generated)], dtype=mx.int32)
+                output = model(
+                    decoder_input_ids=state.canvas,
+                    cache=cache,
+                    previous_confidence=state.confidence,
+                    previous_entropy=state.entropy,
+                    token_age=state.age,
+                    latent_state=state.latent_state,
+                    history_hidden_state=state.history_hidden_state,
+                    return_proposal_statistics=True,
+                    denoise_temperature=denoise_temperature,
+                    repetition_penalty=float(config.repetition_penalty),
+                    compiled_decoder_step=None,
+                    profiler=profiler,
+                )
+                denoise_steps += 1
+                proposal = output.proposal
+                proposal_confidence = output.proposal_confidence
+                token_entropy = output.token_entropy
+                greedy_proposal = output.greedy_proposal
+                greedy_confidence = output.greedy_confidence
+                next_canvas = proposal
+                next_confidence = proposal_confidence.astype(mx.float32)
+                next_entropy = token_entropy.astype(mx.float32)
+                next_latent = replace(
+                    output.next_latent_state,
+                    confidence=next_confidence,
+                    entropy=next_entropy,
+                    age=state.age + 1,
+                    token_changed=(next_canvas != state.canvas).astype(mx.float32),
+                    confidence_delta=next_confidence - state.confidence,
+                    entropy_delta=next_entropy - state.entropy,
+                )
+                next_state = _RollingState(
+                    canvas=next_canvas,
+                    confidence=next_confidence,
+                    entropy=next_entropy,
+                    age=state.age + 1,
+                    latent_state=next_latent,
+                    history_hidden_state=output.heavy_hidden_state,
+                )
+                if profiler is not None:
+                    commit_span = profiler.measure(
+                        "commit",
+                        proposal,
+                        proposal_confidence,
+                        token_entropy,
+                    )
+                policy = select_commit_lengths(
+                    sampled_token_ids=proposal,
+                    normal_failure_rate=fused_commit_failure_rate(
+                        proposal_confidence, token_entropy, vocab_size=vocab_size
+                    ),
+                    previous_failure_rate=fused_commit_failure_rate(
+                        state.confidence, state.entropy, vocab_size=vocab_size
+                    ),
+                    greedy_token_ids=greedy_proposal,
+                    jump_failure_rate=fused_commit_failure_rate(
+                        greedy_confidence, token_entropy, vocab_size=vocab_size
+                    ),
+                    ponder_steps=state.latent_state.ponder_steps,
+                    stagnation_steps=state.latent_state.stagnation_steps,
+                    active_rows=always_active,
+                    remaining_lengths=remaining,
+                    failure_budget=float(config.commit_failure_budget),
+                    jump_failure_budget=float(config.jump_failure_budget),
+                    stop_token_id=stop_token_ids,
+                    max_ponder_steps=int(config.max_ponder_steps),
+                    stagnation_threshold=int(config.jump_on_no_progress_after),
+                    min_progress=float(config.min_trajectory_progress),
+                )
+                if profiler is not None:
+                    commit_span.done(policy.commit_lengths, policy.jump_rows)
+                if profiler is not None:
+                    span = profiler.measure(
+                        "sync",
+                        policy.commit_lengths,
+                        policy.jump_rows,
+                        policy.commit_token_ids,
+                    )
+                mx.eval(
+                    policy.commit_lengths,
+                    policy.jump_rows,
+                    policy.commit_token_ids,
+                )
+                commit_len = int(policy.commit_lengths[0].item())
+                jump = bool(policy.jump_rows[0].item())
+                if profiler is not None:
+                    span.done()
+                next_state = replace(
+                    next_state,
+                    latent_state=replace(
+                        next_state.latent_state,
+                        ponder_steps=policy.ponder_steps,
+                        stagnation_steps=policy.stagnation_steps,
+                    ),
+                )
+                if jump:
+                    next_state = replace(next_state, canvas=policy.commit_token_ids)
+                    jumps += 1
+                committed_ids: list[int] = []
+                if commit_len:
+                    committed_ids = [
+                        int(token)
+                        for token in policy.commit_token_ids[0, :commit_len].tolist()
+                    ]
+                    generated.extend(committed_ids)
+                    committed_block = policy.commit_token_ids[:, :commit_len]
+                    if profiler is not None:
+                        span = profiler.measure("update_cache", committed_block)
+                    cache = model.update_cache(committed_block, cache=cache)
+                    mx.eval(
+                        *[
+                            block.keys
+                            for block in cache
+                            if getattr(block, "keys", None) is not None
+                        ],
+                        *[
+                            block.values
+                            for block in cache
+                            if getattr(block, "values", None) is not None
+                        ],
+                    )
+                    if profiler is not None:
+                        span.done()
+                    shifts += 1
+                state = _shift_state(
+                    next_state,
+                    commit_len,
+                    vocab_size=vocab_size,
+                    canvas_length=canvas_length,
+                    unknown_entropy=unknown_entropy,
+                )
+                if profiler is not None:
+                    profiler.finish_step()
+                if denoise_steps == 1:
+                    first_denoise_seconds = time.perf_counter() - step_started
+                    generate_started = time.perf_counter()
+                if committed_ids:
+                    if turn_end in committed_ids:
+                        stop_reason = "turn_end"
+                        break
+                    if any(
+                        token in stop_token_ids and token != turn_end
+                        for token in committed_ids
+                    ):
+                        stop_reason = "eos"
+                        break
+                if len(generated) >= max_new:
+                    stop_reason = "max_new_tokens"
+                    break
+                if (
+                    max_denoising_steps is not None
+                    and denoise_steps >= max_denoising_steps
+                ):
+                    stop_reason = "max_denoising_steps"
+                    break
+                if denoise_steps >= max_iterations:
+                    stop_reason = "episode_watchdog"
+                    break
+
+    generate_seconds = time.perf_counter() - generate_started
+    if denoise_steps <= 1:
+        generate_seconds = first_denoise_seconds
+        steady_hdps = 0.0
+        overall_hdps = (1.0 / first_denoise_seconds) if first_denoise_seconds else 0.0
+        tokens_per_second = (
+            len(generated) / first_denoise_seconds if first_denoise_seconds else 0.0
+        )
+    else:
+        steady_hdps = (denoise_steps - 1) / max(generate_seconds, 1e-9)
+        overall_hdps = denoise_steps / max(
+            first_denoise_seconds + generate_seconds, 1e-9
+        )
+        tokens_per_second = len(generated) / max(
+            first_denoise_seconds + generate_seconds, 1e-9
+        )
+    average_commit = (len(generated) / shifts) if shifts else 0.0
+    tpf = (len(generated) / denoise_steps) if denoise_steps else 0.0
+    return ModilifyMk1GenerationOutput(
+        sequences=prompt_ids + generated,
+        generated_ids=generated,
+        generated_length=len(generated),
+        denoise_steps=denoise_steps,
+        jump_count=jumps,
+        average_commit_len=average_commit,
+        stop_reason=stop_reason,
+        tokens_per_forward=tpf,
+        prefill_seconds=prefill_seconds,
+        generate_seconds=first_denoise_seconds + generate_seconds,
+        first_denoise_seconds=first_denoise_seconds,
+        heavy_denoise_per_second=overall_hdps,
+        steady_heavy_denoise_per_second=steady_hdps,
+        tokens_per_second=tokens_per_second,
+    )

modilify_mlx/language.py

@@ -0,0 +1,110 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Mk1 wrappers around the borrowed DiffusionGemma trunk layers.
+
+The trunk implementation is imported as an internal dependency. The public
+model type remains ``modilify_mk1``. Two Mk1-specific forwards are installed
+on the constructed trunk:
+
+* Router: softmax over all experts, then top-k and renormalize
+* Canvas merge: RMS-capped latent residual, not previous-logit soft embeds
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_vlm.models.diffusion_gemma.language import (
+    DiffusionGemma4Backbone,
+    Router,
+    geglu,
+)
+
+LATENT_RESIDUAL_RMS_RATIO_CAP = 0.5
+
+
+class Mk1Router(Router):
+    """Official router weights with Mk1 log-softmax / top-k routing."""
+
+    def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]:
+        x = mx.fast.rms_norm(x, None, self.eps)
+        x = x * self.scale * self._root_size
+        scores = self.proj(x)
+        probabilities = mx.softmax(scores, axis=-1, precise=True)
+        top_k = self.config.top_k_experts
+        indices = mx.argpartition(probabilities, kth=-top_k, axis=-1)[..., -top_k:]
+        weights = mx.take_along_axis(probabilities, indices, axis=-1)
+        weights = weights / mx.sum(weights, axis=-1, keepdims=True)
+        weights = weights * self.per_expert_scale[indices]
+        return indices, weights
+
+
+def merge_latent_context(
+    mapper: nn.Module,
+    token_embeddings: mx.array,
+    latent_context: mx.array | None,
+    *,
+    rms_ratio_cap: float = LATENT_RESIDUAL_RMS_RATIO_CAP,
+) -> mx.array:
+    """Apply the native self-conditioning bridge with Mk1 RMS capping."""
+
+    context = (
+        mx.zeros_like(token_embeddings)
+        if latent_context is None
+        else latent_context.astype(token_embeddings.dtype)
+    )
+    if context.shape != token_embeddings.shape:
+        raise ValueError("Latent context must match the canvas embedding shape.")
+    normalized = mapper.pre_norm(context)
+    mapped = mapper.down_proj(
+        geglu(mapper.gate_proj(normalized), mapper.up_proj(normalized))
+    )
+    mapped_rms = mx.sqrt(
+        mx.mean(mx.square(mapped.astype(mx.float32)), axis=-1, keepdims=True)
+    )
+    token_rms = mx.sqrt(
+        mx.mean(
+            mx.square(token_embeddings.astype(mx.float32)), axis=-1, keepdims=True
+        )
+    )
+    cap = rms_ratio_cap * token_rms
+    scale = cap / mx.sqrt(mx.square(mapped_rms) + mx.square(cap) + 1.0e-12)
+    mapped = mapped * scale.astype(mapped.dtype)
+    return mapper.post_norm(token_embeddings + mapped)
+
+
+def _install_mk1_embed_canvas(decoder: nn.Module) -> None:
+    def _embed_canvas(
+        canvas_ids,
+        self_conditioning_logits=None,
+        self_conditioning_embeddings=None,
+    ):
+        if self_conditioning_logits is not None:
+            raise ValueError(
+                "Modilify Mk1 uses latent embeddings, not logits self-conditioning."
+            )
+        token_embeddings = decoder.embed_tokens(canvas_ids) * decoder.embed_scale
+        return merge_latent_context(
+            decoder.self_conditioning,
+            token_embeddings,
+            self_conditioning_embeddings,
+        )
+
+    decoder._embed_canvas = _embed_canvas
+
+
+def _install_mk1_routers(backbone: DiffusionGemma4Backbone) -> None:
+    for layer in backbone.decoder.layers:
+        replacement = Mk1Router(layer.router.config)
+        replacement.update(layer.router.parameters())
+        layer.router = replacement
+
+
+def build_mk1_backbone(trunk_config) -> DiffusionGemma4Backbone:
+    """Construct the borrowed trunk and install Mk1 forwards."""
+
+    backbone = DiffusionGemma4Backbone(trunk_config)
+    _install_mk1_routers(backbone)
+    _install_mk1_embed_canvas(backbone.decoder)
+    return backbone

modilify_mlx/latent_deliberation.py

@@ -0,0 +1,449 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Fixed-shape latent deliberation state for Modilify Mk1 decoding.
+
+PyTorch → MLX:
+    nn.LayerNorm                     mlx.nn.LayerNorm
+    nn.Linear                        mlx.nn.Linear
+    nn.MultiheadAttention in_proj    BiasedMHA query/key/value_proj (split QKV)
+    nn.Sequential token_ff/memory_ff mlx.nn.Sequential (layers.0 / layers.2)
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+@dataclass
+class LatentDeliberationState:
+    """Persistent, fixed-size state for one or more canvas episodes."""
+
+    token_latents: mx.array
+    memory_slots: mx.array
+    confidence: mx.array
+    entropy: mx.array
+    age: mx.array
+    token_changed: mx.array
+    confidence_delta: mx.array
+    entropy_delta: mx.array
+    ponder_steps: mx.array
+    stagnation_steps: mx.array
+
+    @classmethod
+    def empty(
+        cls,
+        *,
+        batch_size: int,
+        canvas_length: int,
+        latent_dim: int,
+        memory_slots: int,
+        dtype: mx.Dtype,
+    ) -> "LatentDeliberationState":
+        """Create a zero-initialized recurrent state."""
+
+        return cls(
+            token_latents=mx.zeros((batch_size, canvas_length, latent_dim), dtype=dtype),
+            memory_slots=mx.zeros((batch_size, memory_slots, latent_dim), dtype=dtype),
+            confidence=mx.zeros((batch_size, canvas_length), dtype=mx.float32),
+            entropy=mx.zeros((batch_size, canvas_length), dtype=mx.float32),
+            age=mx.zeros((batch_size, canvas_length), dtype=mx.int32),
+            token_changed=mx.zeros((batch_size, canvas_length), dtype=mx.float32),
+            confidence_delta=mx.zeros((batch_size, canvas_length), dtype=mx.float32),
+            entropy_delta=mx.zeros((batch_size, canvas_length), dtype=mx.float32),
+            ponder_steps=mx.zeros((batch_size,), dtype=mx.int32),
+            stagnation_steps=mx.zeros((batch_size,), dtype=mx.int32),
+        )
+
+    def shift(
+        self, committed: int, *, entropy_fill_value: float = 0.0
+    ) -> "LatentDeliberationState":
+        """Drop committed canvas positions without changing long-term memory."""
+
+        canvas_length = self.token_latents.shape[1]
+        if not 0 <= committed <= canvas_length:
+            raise ValueError("`committed` must be in [0, canvas_length].")
+        if committed == 0:
+            return self
+
+        def shifted(tensor: mx.array, fill_value: float | int = 0) -> mx.array:
+            result = mx.full(tensor.shape, fill_value, dtype=tensor.dtype)
+            if committed < canvas_length:
+                kept = tensor[:, committed:]
+                result = mx.concatenate(
+                    [kept, mx.full(
+                        (tensor.shape[0], committed, *tensor.shape[2:]),
+                        fill_value,
+                        dtype=tensor.dtype,
+                    )],
+                    axis=1,
+                ) if tensor.ndim > 2 else mx.concatenate(
+                    [
+                        kept,
+                        mx.full(
+                            (tensor.shape[0], committed),
+                            fill_value,
+                            dtype=tensor.dtype,
+                        ),
+                    ],
+                    axis=1,
+                )
+            return result
+
+        return LatentDeliberationState(
+            token_latents=shifted(self.token_latents),
+            memory_slots=self.memory_slots,
+            confidence=shifted(self.confidence),
+            entropy=shifted(self.entropy, entropy_fill_value),
+            age=shifted(self.age),
+            token_changed=shifted(self.token_changed),
+            confidence_delta=shifted(self.confidence_delta),
+            entropy_delta=shifted(self.entropy_delta),
+            ponder_steps=mx.zeros_like(self.ponder_steps),
+            stagnation_steps=mx.zeros_like(self.stagnation_steps),
+        )
+
+
+def advance_trajectory_clocks(
+    ponder_steps: mx.array,
+    stagnation_steps: mx.array,
+    *,
+    commit_lengths: mx.array,
+    active_rows: mx.array,
+    progress_scores: mx.array,
+    min_progress: float,
+) -> tuple[mx.array, mx.array]:
+    """Advance useful-ponder and true-stagnation clocks for each row."""
+
+    if min_progress < 0:
+        raise ValueError("`min_progress` must be non-negative.")
+    if not (
+        ponder_steps.shape
+        == stagnation_steps.shape
+        == commit_lengths.shape
+        == active_rows.shape
+        == progress_scores.shape
+    ):
+        raise ValueError("Trajectory clock inputs must share shape [batch].")
+    committed = commit_lengths > 0
+    waiting = active_rows & (~committed)
+    improving = progress_scores >= min_progress
+    next_ponder = mx.where(
+        committed,
+        mx.zeros_like(ponder_steps),
+        ponder_steps + waiting.astype(mx.int32),
+    )
+    next_stagnation = mx.where(
+        committed,
+        mx.zeros_like(stagnation_steps),
+        mx.where(
+            waiting & improving,
+            mx.zeros_like(stagnation_steps),
+            stagnation_steps + waiting.astype(mx.int32),
+        ),
+    )
+    return next_ponder.astype(mx.int32), next_stagnation.astype(mx.int32)
+
+
+def should_force_trajectory_jump(
+    ponder_steps: mx.array,
+    stagnation_steps: mx.array,
+    *,
+    max_ponder_steps: int,
+    stagnation_threshold: int,
+) -> mx.array:
+    """Return whether a row has exhausted ponder or stagnation budget."""
+
+    if max_ponder_steps <= 0 or stagnation_threshold <= 0:
+        raise ValueError("Trajectory jump limits must be positive.")
+    return (ponder_steps >= max_ponder_steps) | (
+        stagnation_steps >= stagnation_threshold
+    )
+
+
+class BiasedMHA(nn.Module):
+    """Explicit QKV attention matching PyTorch MultiheadAttention weights."""
+
+    def __init__(self, dims: int, num_heads: int) -> None:
+        super().__init__()
+        if dims % num_heads:
+            raise ValueError("`dims` must be divisible by `num_heads`.")
+        self.num_heads = num_heads
+        self.head_dim = dims // num_heads
+        self.scale = self.head_dim**-0.5
+        self.query_proj = nn.Linear(dims, dims, bias=True)
+        self.key_proj = nn.Linear(dims, dims, bias=True)
+        self.value_proj = nn.Linear(dims, dims, bias=True)
+        self.out_proj = nn.Linear(dims, dims, bias=True)
+
+    def __call__(
+        self,
+        queries: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        mask: mx.array | None = None,
+    ) -> mx.array:
+        batch, query_len, dims = queries.shape
+        key_len = keys.shape[1]
+        queries = self.query_proj(queries).reshape(
+            batch, query_len, self.num_heads, self.head_dim
+        ).transpose(0, 2, 1, 3)
+        keys = self.key_proj(keys).reshape(
+            batch, key_len, self.num_heads, self.head_dim
+        ).transpose(0, 2, 1, 3)
+        values = self.value_proj(values).reshape(
+            batch, key_len, self.num_heads, self.head_dim
+        ).transpose(0, 2, 1, 3)
+        attn_mask = None
+        if mask is not None:
+            attn_mask = mask.reshape(1, 1, mask.shape[-2], mask.shape[-1])
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=self.scale, mask=attn_mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(batch, query_len, dims)
+        return self.out_proj(output)
+
+
+def _silu(x: mx.array) -> mx.array:
+    return x * mx.sigmoid(x)
+
+
+class _TemporalTransformerCell(nn.Module):
+    """One-step recurrent token update with fixed-slot memory attention."""
+
+    def __init__(
+        self,
+        latent_dim: int,
+        num_heads: int,
+        local_attention_window: int,
+    ) -> None:
+        super().__init__()
+        self.state_norm = nn.LayerNorm(latent_dim)
+        self.observation_norm = nn.LayerNorm(latent_dim)
+        self.memory_address_norm = nn.LayerNorm(latent_dim)
+        self.memory_value_norm = nn.LayerNorm(latent_dim)
+        self.temporal_update = nn.Linear(2 * latent_dim, 2 * latent_dim)
+        self.local_attention = BiasedMHA(latent_dim, num_heads)
+        self.local_attention_window = local_attention_window
+        self.token_memory_attention = BiasedMHA(latent_dim, num_heads)
+        self.memory_token_attention = BiasedMHA(latent_dim, num_heads)
+        self.token_ff_norm = nn.LayerNorm(latent_dim)
+        self.memory_ff_norm = nn.LayerNorm(latent_dim)
+        self.stored_token_norm = nn.LayerNorm(latent_dim)
+        self.stored_memory_norm = nn.LayerNorm(latent_dim)
+        expansion = latent_dim * 4
+        self.token_ff = nn.Sequential(
+            nn.Linear(latent_dim, expansion),
+            nn.SiLU(),
+            nn.Linear(expansion, latent_dim),
+        )
+        self.memory_ff = nn.Sequential(
+            nn.Linear(latent_dim, expansion),
+            nn.SiLU(),
+            nn.Linear(expansion, latent_dim),
+        )
+
+    def _local_mask(self, length: int, dtype: mx.Dtype) -> mx.array:
+        positions = mx.arange(length)
+        allowed = (
+            mx.abs(positions[:, None] - positions[None, :]) < self.local_attention_window
+        )
+        return mx.where(
+            allowed,
+            mx.zeros((length, length), dtype=dtype),
+            mx.full((length, length), mx.finfo(dtype).min, dtype=dtype),
+        )
+
+    def __call__(
+        self,
+        previous_tokens: mx.array,
+        observation: mx.array,
+        memory: mx.array,
+        memory_slot_identity: mx.array,
+    ) -> tuple[mx.array, mx.array]:
+        gate_logits, candidate = mx.split(
+            self.temporal_update(
+                mx.concatenate(
+                    (
+                        self.state_norm(previous_tokens),
+                        self.observation_norm(observation),
+                    ),
+                    axis=-1,
+                )
+            ),
+            2,
+            axis=-1,
+        )
+        gate = mx.sigmoid(gate_logits)
+        tokens = gate * previous_tokens + (1.0 - gate) * _silu(candidate)
+        local_mask = self._local_mask(tokens.shape[1], tokens.dtype)
+        normed_tokens = self.state_norm(tokens)
+        tokens = tokens + self.local_attention(
+            normed_tokens, normed_tokens, normed_tokens, mask=local_mask
+        )
+
+        addressed_memory = self.memory_address_norm(memory + memory_slot_identity)
+        memory_values = self.memory_value_norm(memory)
+        normed_tokens = self.state_norm(tokens)
+        tokens = tokens + self.token_memory_attention(
+            normed_tokens, addressed_memory, memory_values
+        )
+        tokens = tokens + self.token_ff(self.token_ff_norm(tokens))
+
+        normed_tokens = self.state_norm(tokens)
+        memory = memory + self.memory_token_attention(
+            addressed_memory, normed_tokens, normed_tokens
+        )
+        memory = memory + self.memory_ff(self.memory_ff_norm(memory))
+        return (
+            self.stored_token_norm(tokens),
+            self.stored_memory_norm(memory) + memory_slot_identity,
+        )
+
+
+class LatentDeliberationTransformer(nn.Module):
+    """Small recurrent Transformer that compresses repeated denoise context."""
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        latent_dim: int = 512,
+        memory_slots: int = 16,
+        num_layers: int = 2,
+        num_heads: int = 8,
+        local_attention_window: int = 32,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        del dropout
+        if latent_dim % num_heads:
+            raise ValueError("`latent_dim` must be divisible by `num_heads`.")
+        if local_attention_window <= 0:
+            raise ValueError("`local_attention_window` must be positive.")
+        self.hidden_size = hidden_size
+        self.latent_dim = latent_dim
+        self.memory_slots = memory_slots
+        self.heavy_projection = nn.Linear(hidden_size, latent_dim, bias=False)
+        self.embedding_projection = nn.Linear(hidden_size, latent_dim, bias=False)
+        self.scalar_projection = nn.Linear(11, latent_dim, bias=False)
+        self.blocks = [
+            _TemporalTransformerCell(latent_dim, num_heads, local_attention_window)
+            for _ in range(num_layers)
+        ]
+        self.output_norm = nn.LayerNorm(latent_dim)
+        self.output_projection = nn.Linear(latent_dim, hidden_size, bias=False)
+        self.memory_slot_identity = mx.zeros((memory_slots, latent_dim))
+
+    def scaled_memory_slot_identity(
+        self,
+        *,
+        batch_size: int,
+        dtype: mx.Dtype,
+    ) -> mx.array:
+        """Return unit directions scaled to LayerNorm RMS."""
+
+        identity = self.memory_slot_identity.astype(mx.float32)
+        identity = identity / mx.maximum(
+            mx.linalg.norm(identity, axis=-1, keepdims=True), 1.0e-12
+        )
+        identity = identity * math.sqrt(self.latent_dim)
+        return mx.broadcast_to(
+            identity.astype(dtype)[None, :, :],
+            (batch_size, self.memory_slots, self.latent_dim),
+        )
+
+    def project_context(self, token_latents: mx.array) -> mx.array:
+        """Translate latent state into a self-conditioning embedding."""
+
+        return self.output_projection(self.output_norm(token_latents))
+
+    def __call__(
+        self,
+        *,
+        heavy_hidden: mx.array,
+        token_embeddings: mx.array,
+        confidence: mx.array,
+        entropy: mx.array,
+        state: LatentDeliberationState,
+    ) -> tuple[mx.array, LatentDeliberationState]:
+        """Advance latent memory and produce decoder self-conditioning."""
+
+        if heavy_hidden.ndim != 3:
+            raise ValueError("`heavy_hidden` must have shape [batch, canvas, hidden].")
+        if heavy_hidden.shape != token_embeddings.shape:
+            raise ValueError(
+                "`heavy_hidden` and `token_embeddings` must have the same shape."
+            )
+        batch_size, canvas_length, hidden_size = heavy_hidden.shape
+        if hidden_size != self.hidden_size:
+            raise ValueError("Unexpected hidden size for latent deliberation.")
+        expected_state = (batch_size, canvas_length, self.latent_dim)
+        if state.token_latents.shape != expected_state:
+            raise ValueError("State token latents do not match the current canvas.")
+        if state.memory_slots.shape != (
+            batch_size,
+            self.memory_slots,
+            self.latent_dim,
+        ):
+            raise ValueError("State memory slots do not match this module.")
+
+        dtype = heavy_hidden.dtype
+        position = mx.linspace(-1.0, 1.0, canvas_length, dtype=dtype)
+        position = mx.broadcast_to(position[None, :], (batch_size, canvas_length))
+        age = mx.minimum(state.age.astype(dtype), 32767).astype(dtype)
+        entropy_delta = state.entropy_delta.astype(dtype)
+        scalars = mx.stack(
+            (
+                confidence.astype(dtype),
+                mx.log1p(entropy.astype(dtype)),
+                mx.log1p(age),
+                position,
+                state.token_changed.astype(dtype),
+                state.confidence_delta.astype(dtype),
+                mx.sign(entropy_delta) * mx.log1p(mx.abs(entropy_delta)),
+                mx.broadcast_to(
+                    mx.log1p(state.ponder_steps.astype(dtype))[:, None],
+                    (batch_size, canvas_length),
+                ),
+                mx.broadcast_to(
+                    mx.log1p(state.stagnation_steps.astype(dtype))[:, None],
+                    (batch_size, canvas_length),
+                ),
+                confidence.astype(dtype)
+                * mx.exp(-mx.maximum(entropy.astype(dtype), 0.0)),
+                mx.maximum(state.confidence_delta.astype(dtype), 0.0)
+                + mx.log1p(mx.maximum(-entropy_delta, 0.0)),
+            ),
+            axis=-1,
+        )
+        observation = (
+            self.heavy_projection(heavy_hidden)
+            + self.embedding_projection(token_embeddings)
+            + self.scalar_projection(scalars)
+        )
+        tokens = state.token_latents
+        memory = state.memory_slots
+        slot_identity = self.scaled_memory_slot_identity(
+            batch_size=batch_size,
+            dtype=memory.dtype,
+        )
+        for block in self.blocks:
+            tokens, memory = block(tokens, observation, memory, slot_identity)
+            observation = tokens
+        next_state = LatentDeliberationState(
+            token_latents=tokens,
+            memory_slots=memory,
+            confidence=confidence.astype(mx.float32),
+            entropy=entropy.astype(mx.float32),
+            age=state.age,
+            token_changed=state.token_changed,
+            confidence_delta=state.confidence_delta,
+            entropy_delta=state.entropy_delta,
+            ponder_steps=state.ponder_steps,
+            stagnation_steps=state.stagnation_steps,
+        )
+        return self.project_context(tokens), next_state

modilify_mlx/modeling.py

@@ -0,0 +1,433 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Native MLX Modilify Mk1 model."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten
+
+from .config import MODEL_TYPE, ModilifyMk1Config
+from .convert_utils import remap_state_dict
+from .fast_decode import (
+    decoder_hidden_states,
+    make_compiled_attn_layers,
+    quantize_experts,
+)
+from .language import build_mk1_backbone
+from .latent_deliberation import LatentDeliberationState, LatentDeliberationTransformer
+
+
+@dataclass
+class ModilifyMk1StepOutput:
+    """One heavy-denoise step over the rolling canvas."""
+
+    logits: mx.array | None
+    heavy_hidden_state: mx.array
+    next_latent_state: LatentDeliberationState
+    cache: Any
+    latent_context: mx.array
+    proposal: mx.array | None = None
+    proposal_confidence: mx.array | None = None
+    token_entropy: mx.array | None = None
+    greedy_proposal: mx.array | None = None
+    greedy_confidence: mx.array | None = None
+
+
+def _softcap(logits: mx.array, cap: float) -> mx.array:
+    return mx.tanh(logits.astype(mx.float32) / cap) * cap
+
+
+def _softmax_statistics(
+    logits: mx.array,
+    temperature: float,
+) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]:
+    scores = logits.astype(mx.float32) / temperature
+    probabilities = mx.softmax(scores, axis=-1, precise=True)
+    greedy_proposal = mx.argmax(probabilities, axis=-1)
+    greedy_confidence = mx.squeeze(
+        mx.take_along_axis(probabilities, greedy_proposal[..., None], axis=-1),
+        axis=-1,
+    )
+    token_entropy = -mx.sum(
+        probabilities * mx.log(mx.maximum(probabilities, 1.0e-30)),
+        axis=-1,
+    )
+    return scores, probabilities, greedy_proposal, greedy_confidence, token_entropy
+
+
+_compiled_softmax_statistics = mx.compile(_softmax_statistics, shapeless=True)
+
+
+class ModilifyMk1ForBlockDiffusion(nn.Module):
+    """Inference-only multimodal Modilify Mk1 model."""
+
+    def __init__(self, config: ModilifyMk1Config) -> None:
+        super().__init__()
+        if config.model_type != MODEL_TYPE:
+            raise ValueError(
+                f"Refusing to construct Mk1 with model_type={config.model_type!r}."
+            )
+        self.config = config
+        self.model = build_mk1_backbone(config.trunk_model_config())
+        self.latent_deliberation = LatentDeliberationTransformer(
+            hidden_size=config.hidden_size,
+            latent_dim=config.latent_dim,
+            memory_slots=config.latent_memory_slots,
+            num_layers=config.latent_num_layers,
+            num_heads=config.latent_num_heads,
+            local_attention_window=config.latent_local_attention_window,
+            dropout=config.latent_dropout,
+        )
+        self.final_logit_softcapping = float(
+            config.text_config.final_logit_softcapping
+        )
+        self._decoder_compile_failed = False
+        self._compiled_attn_layers = None
+
+    def make_cache(self, max_size: int | None = None):
+        return self.model.encoder.make_cache(max_size=max_size)
+
+    def embed_canvas_tokens(self, decoder_input_ids: mx.array) -> mx.array:
+        return (
+            self.model.decoder.embed_tokens(decoder_input_ids)
+            * self.model.decoder.embed_scale
+        )
+
+    def prefill(
+        self,
+        input_ids: mx.array,
+        *,
+        attention_mask: mx.array | None = None,
+        cache=None,
+        pixel_values: mx.array | None = None,
+        mm_token_type_ids: mx.array | None = None,
+    ):
+        if cache is None:
+            cache = self.make_cache()
+        _, cache = self.model.encoder(
+            input_ids,
+            attention_mask=attention_mask,
+            cache=cache,
+            pixel_values=pixel_values,
+            mm_token_type_ids=mm_token_type_ids,
+        )
+        return cache
+
+    def update_cache(self, input_ids: mx.array, *, cache, attention_mask=None):
+        _, cache = self.model.encoder(
+            input_ids,
+            attention_mask=attention_mask,
+            cache=cache,
+        )
+        return cache
+
+    def _prepare_latent_context(
+        self,
+        decoder_input_ids: mx.array,
+        *,
+        history_hidden_state: mx.array | None,
+        confidence: mx.array | None,
+        entropy: mx.array | None,
+        age: mx.array | None,
+        latent_state: LatentDeliberationState | None,
+        dtype: mx.Dtype,
+    ) -> tuple[mx.array, LatentDeliberationState]:
+        batch_size, canvas_length = decoder_input_ids.shape
+        if latent_state is None:
+            latent_state = LatentDeliberationState.empty(
+                batch_size=batch_size,
+                canvas_length=canvas_length,
+                latent_dim=self.config.latent_dim,
+                memory_slots=self.config.latent_memory_slots,
+                dtype=dtype,
+            )
+        if confidence is None:
+            confidence = latent_state.confidence
+        else:
+            confidence = mx.squeeze(confidence.astype(mx.float32), axis=-1) if (
+                confidence.ndim == 3
+            ) else confidence.astype(mx.float32)
+        if entropy is None:
+            entropy = latent_state.entropy
+        else:
+            entropy = mx.squeeze(entropy.astype(mx.float32), axis=-1) if (
+                entropy.ndim == 3
+            ) else entropy.astype(mx.float32)
+        if age is not None:
+            latent_state = replace(latent_state, age=age.astype(mx.int32))
+        token_embeddings = self.embed_canvas_tokens(decoder_input_ids)
+        history = (
+            mx.zeros_like(token_embeddings)
+            if history_hidden_state is None
+            else history_hidden_state
+        )
+        return self.latent_deliberation(
+            heavy_hidden=history,
+            token_embeddings=token_embeddings,
+            confidence=confidence,
+            entropy=entropy,
+            state=latent_state,
+        )
+
+    def _proposal_statistics(
+        self,
+        logits: mx.array,
+        *,
+        denoise_temperature: float | None = None,
+        repetition_token_mask: mx.array | None = None,
+        repetition_penalty: float = 1.0,
+    ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]:
+        temperature = (
+            self.config.denoise_temperature
+            if denoise_temperature is None
+            else float(denoise_temperature)
+        )
+        if temperature <= 0:
+            raise ValueError("`denoise_temperature` must be positive.")
+        if (
+            repetition_token_mask is not None
+            and repetition_penalty != 1.0
+            and repetition_penalty > 0
+        ):
+            scores = logits.astype(mx.float32)
+            penalized = mx.where(
+                scores < 0,
+                scores * repetition_penalty,
+                scores / repetition_penalty,
+            )
+            mask = repetition_token_mask.astype(mx.bool_)[:, None, :]
+            logits = mx.where(mask, penalized, scores)
+        try:
+            scores, probabilities, greedy_proposal, greedy_confidence, token_entropy = (
+                _compiled_softmax_statistics(logits, temperature)
+            )
+        except ValueError:
+            scores, probabilities, greedy_proposal, greedy_confidence, token_entropy = (
+                _softmax_statistics(logits, temperature)
+            )
+        proposal = mx.random.categorical(scores, axis=-1)
+        proposal_confidence = mx.squeeze(
+            mx.take_along_axis(probabilities, proposal[..., None], axis=-1),
+            axis=-1,
+        )
+        return (
+            proposal,
+            proposal_confidence,
+            token_entropy,
+            greedy_proposal,
+            greedy_confidence,
+        )
+
+    def compile_attention(self, cache) -> None:
+        """Compile per-layer attention residuals. Expert FFNs stay eager."""
+
+        from .fast_decode import _cache_capacity, build_decoder_masks
+
+        print("[mk1] compiling attention layers", flush=True)
+        compiled = make_compiled_attn_layers(self.model.decoder, cache)
+        prefix_len = int(getattr(cache[0], "offset", 0))
+        canvas = int(self.config.canvas_length)
+        hidden = int(self.config.hidden_size)
+        dtype = self.model.decoder.embed_tokens.weight.dtype
+        dummy = mx.zeros((1, canvas, hidden), dtype=dtype)
+        offset = mx.array(prefix_len)
+        full_mask, slide_mask = build_decoder_masks(
+            prefix_len=prefix_len,
+            canvas_length=canvas,
+            cache_capacity=max(_cache_capacity(cache), 1),
+            sliding_window=int(self.config.text_config.sliding_window),
+        )
+        try:
+            for layer, attn_fn in zip(self.model.decoder.layers, compiled):
+                mask = (
+                    slide_mask
+                    if layer.layer_type == "sliding_attention"
+                    else full_mask
+                )
+                dummy = attn_fn(dummy, offset, mask)
+            mx.eval(dummy)
+            self._compiled_attn_layers = compiled
+            print("[mk1] attention compile ready", flush=True)
+        except ValueError as exc:
+            print(f"[mk1] attention compile fallback: {exc}", flush=True)
+            self._compiled_attn_layers = None
+
+    def decoder_logits(
+        self,
+        decoder_input_ids: mx.array,
+        latent_context: mx.array,
+        cache,
+        offset: mx.array,
+        full_mask: mx.array,
+        slide_mask: mx.array,
+        compiled_decoder_step=None,
+    ) -> tuple[mx.array, mx.array]:
+        del compiled_decoder_step
+        hidden_states = decoder_hidden_states(
+            self.model.decoder,
+            decoder_input_ids,
+            latent_context,
+            cache,
+            offset,
+            full_mask,
+            slide_mask,
+            compiled_attn_layers=self._compiled_attn_layers,
+        )
+        logits = self.model.decoder.embed_tokens.as_linear(hidden_states)
+        return _softcap(logits, self.final_logit_softcapping), hidden_states
+
+    def __call__(
+        self,
+        *,
+        decoder_input_ids: mx.array,
+        cache,
+        previous_confidence: mx.array | None = None,
+        previous_entropy: mx.array | None = None,
+        token_age: mx.array | None = None,
+        latent_state: LatentDeliberationState | None = None,
+        history_hidden_state: mx.array | None = None,
+        decoder_attention_mask: mx.array | None = None,
+        return_proposal_statistics: bool = False,
+        denoise_temperature: float | None = None,
+        repetition_token_mask: mx.array | None = None,
+        repetition_penalty: float = 1.0,
+        compiled_decoder_step=None,
+        profiler=None,
+    ) -> ModilifyMk1StepOutput:
+        """Run one inference step over a noisy diffusion canvas."""
+
+        del compiled_decoder_step
+        dtype = self.model.decoder.embed_tokens.weight.dtype
+        if profiler is not None:
+            span = profiler.measure("latent", decoder_input_ids)
+        latent_context, next_state = self._prepare_latent_context(
+            decoder_input_ids,
+            history_hidden_state=history_hidden_state,
+            confidence=previous_confidence,
+            entropy=previous_entropy,
+            age=token_age,
+            latent_state=latent_state,
+            dtype=dtype,
+        )
+        if profiler is not None:
+            span.done(latent_context, next_state.token_latents, next_state.memory_slots)
+        del decoder_attention_mask
+        from .fast_decode import _cache_capacity, build_decoder_masks
+
+        prefix_len = int(getattr(cache[0], "offset", 0))
+        offset = mx.array(prefix_len)
+        canvas_length = int(decoder_input_ids.shape[1])
+        full_mask, slide_mask = build_decoder_masks(
+            prefix_len=prefix_len,
+            canvas_length=canvas_length,
+            cache_capacity=max(_cache_capacity(cache), 1),
+            sliding_window=int(self.config.text_config.sliding_window),
+            batch_size=int(decoder_input_ids.shape[0]),
+        )
+        hidden_states = decoder_hidden_states(
+            self.model.decoder,
+            decoder_input_ids,
+            latent_context,
+            cache,
+            offset,
+            full_mask,
+            slide_mask,
+            compiled_attn_layers=None,
+            profiler=profiler,
+        )
+        if profiler is not None:
+            span = profiler.measure("lm_head", hidden_states)
+        logits = self.model.decoder.embed_tokens.as_linear(hidden_states)
+        logits = _softcap(logits, self.final_logit_softcapping)
+        if profiler is not None:
+            span.done(logits)
+        statistics = (None, None, None, None, None)
+        if return_proposal_statistics:
+            if profiler is not None:
+                span = profiler.measure("softmax", logits)
+            statistics = self._proposal_statistics(
+                logits,
+                denoise_temperature=denoise_temperature,
+                repetition_token_mask=repetition_token_mask,
+                repetition_penalty=repetition_penalty,
+            )
+            if profiler is not None:
+                span.done(
+                    statistics[0],
+                    statistics[1],
+                    statistics[2],
+                    statistics[3],
+                    statistics[4],
+                )
+        return ModilifyMk1StepOutput(
+            logits=None if return_proposal_statistics else logits,
+            heavy_hidden_state=hidden_states,
+            next_latent_state=next_state,
+            cache=cache,
+            latent_context=latent_context,
+            proposal=statistics[0],
+            proposal_confidence=statistics[1],
+            token_entropy=statistics[2],
+            greedy_proposal=statistics[3],
+            greedy_confidence=statistics[4],
+        )
+
+
+def _load_weight_files(model_path: Path) -> dict[str, mx.array]:
+    weight_files = sorted(model_path.glob("*.safetensors"))
+    if not weight_files:
+        raise FileNotFoundError(f"No safetensors found in {model_path}")
+    weights: dict[str, mx.array] = {}
+    for weight_file in weight_files:
+        weights.update(mx.load(str(weight_file)))
+    return weights
+
+
+def load(
+    model_path: str | Path,
+    *,
+    lazy: bool = False,
+    expert_bits: int = 16,
+    expert_group_size: int = 64,
+):
+    """Load a native ``modilify_mk1`` MLX checkpoint."""
+
+    model_path = Path(model_path)
+    config = ModilifyMk1Config.from_json(model_path / "config.json")
+    if config.model_type != MODEL_TYPE:
+        raise ValueError(
+            f"Refusing to load model_type={config.model_type!r}; "
+            f"expected {MODEL_TYPE!r}."
+        )
+    print("[mk1] constructing graph", flush=True)
+    model = ModilifyMk1ForBlockDiffusion(config)
+    print("[mk1] reading shards", flush=True)
+    weights = _load_weight_files(model_path)
+    remapped = remap_state_dict(weights.items())
+    if len(remapped) != len(weights) or any(key not in remapped for key in weights):
+        print(
+            f"[mk1] remapped {len(weights)} source tensors -> {len(remapped)} MLX tensors",
+            flush=True,
+        )
+    del weights
+    print(f"[mk1] loading {len(remapped)} tensors", flush=True)
+    model.load_weights(list(remapped.items()), strict=True)
+    del remapped
+    if not lazy:
+        print("[mk1] evaluating parameters", flush=True)
+        mx.eval(model.parameters())
+    if expert_bits and expert_bits < 16:
+        quantize_experts(
+            model, bits=int(expert_bits), group_size=int(expert_group_size)
+        )
+    return model, config
+
+
+def parameter_names(model: nn.Module) -> list[str]:
+    return [name for name, _ in tree_flatten(model.parameters())]

modilify_mlx/profile.py

@@ -0,0 +1,64 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Optional per-denoise phase timers. Enabled only with --profile."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+import time
+
+import mlx.core as mx
+
+PHASES = (
+    "latent",
+    "attn",
+    "moe",
+    "lm_head",
+    "softmax",
+    "commit",
+    "sync",
+    "update_cache",
+)
+
+
+class DenoiseProfiler:
+    def __init__(self) -> None:
+        self.totals = defaultdict(float)
+        self.steps = 0
+
+    def add(self, phase: str, seconds: float) -> None:
+        self.totals[phase] += float(seconds)
+
+    def finish_step(self) -> None:
+        self.steps += 1
+
+    def measure(self, phase: str, *arrays: mx.array):
+        mx.eval(*arrays)
+        started = time.perf_counter()
+
+        class _Span:
+            def __init__(self, profiler: DenoiseProfiler, name: str) -> None:
+                self.profiler = profiler
+                self.name = name
+                self.started = started
+
+            def done(self, *outputs: mx.array) -> None:
+                if outputs:
+                    mx.eval(*outputs)
+                self.profiler.add(self.name, time.perf_counter() - self.started)
+
+        return _Span(self, phase)
+
+    def summary(self) -> str:
+        counted = sum(self.totals[name] for name in PHASES)
+        lines = [
+            f"[profile] steps={self.steps} accounted={counted:.3f}s",
+        ]
+        for name in PHASES:
+            value = self.totals[name]
+            share = (100.0 * value / counted) if counted else 0.0
+            per = (value / self.steps) if self.steps else 0.0
+            lines.append(
+                f"[profile] {name:12s}  {value:7.3f}s  {share:5.1f}%  {per*1000:6.1f} ms/step"
+            )
+        return "\n".join(lines)