wisp_mtp_reference.py

#!/usr/bin/env python3
"""Strict correctness runtime for a Wisp Hugging Face export.

The runtime consumes four model files from one complete, manifest-bound export:

    config.json
    model.safetensors
    mtp_config.json
    mtp.safetensors

``export_manifest.json`` and its complete declared payload are required so the
four consumed files cannot be mixed across exports without detection.
``model.safetensors`` uses ordinary Hugging Face Llama parameter names.
``mtp.safetensors`` uses Wisp's native ``mtp.*`` names.  The loader maps the
trunk into the packaged inference-only Wisp implementation so the MTP module sees the
raw, pre-final-norm trunk residual that it saw during training.  A generic
``LlamaModel.last_hidden_state`` is post-final-norm and is therefore not a
compatible substitute.

This module is a fail-closed correctness oracle.  Its MTP route recomputes full
prefixes and has no rollback-capable target KV cache.  It reports route
telemetry, but intentionally reports no latency or throughput number and makes
no production-latency claim.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
import hashlib
import importlib.util
import json
import math
import os
import secrets
import stat
import sys
from typing import Any

import mlx.core as mx
from mlx.utils import tree_flatten, tree_unflatten


REQUIRED_PACKAGE_FILES = (
    "config.json",
    "model.safetensors",
    "mtp_config.json",
    "mtp.safetensors",
    "export_manifest.json",
)

EXPECTED_RELEASE_PAYLOAD_FILES = (
    "LICENSE",
    "README.md",
    "config.json",
    "generation_config.json",
    "model.safetensors",
    "mtp.safetensors",
    "mtp_config.json",
    "special_tokens_map.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "wisp_mtp_model.py",
    "wisp_mtp_reference.py",
)

EXPORT_MANIFEST_KEYS = frozenset(
    {
        "schema_version",
        "repo_id",
        "release_complete",
        "portable_evidence_bundle_required",
        "evaluation_sources",
        "model_card_template_sha256",
        "source_checkpoint",
        "trunk_parameters",
        "mtp_parameters_excluding_shared_embedding_and_head",
        "files",
    }
)

EVALUATION_SOURCE_KEYS = frozenset(
    {
        "validation",
        "acceptance_comparison",
        "format_ablation",
        "rollout",
    }
)

MODEL_CONFIG_KEYS = frozenset(
    {
        "architectures",
        "model_type",
        "hidden_size",
        "intermediate_size",
        "num_hidden_layers",
        "num_attention_heads",
        "num_key_value_heads",
        "head_dim",
        "max_position_embeddings",
        "rms_norm_eps",
        "rope_theta",
        "vocab_size",
        "tie_word_embeddings",
        "hidden_act",
        "attention_bias",
        "mlp_bias",
        "torch_dtype",
        "bos_token_id",
        "eos_token_id",
        "pad_token_id",
    }
)

MTP_CONFIG_KEYS = frozenset(
    {
        "mtp_layers",
        "mtp_depth_trained",
        "shared_lm_head",
        "recursive",
        "note",
        "trained_steps",
        "schema_version",
        "architecture",
        "hidden_state_stage",
        "requires_full_sequence_attention",
        "tensor_prefix",
    }
)

MTP_SCHEMA_VERSION = 1
MTP_ARCHITECTURE = "wisp_recursive_shared_module"
MTP_HIDDEN_STATE_STAGE = "trunk_pre_final_norm_residual"
MTP_REQUIRES_FULL_SEQUENCE_ATTENTION = True
MTP_TENSOR_PREFIX = "mtp."

EXPECTED_MTP_NOTE = (
    "One shared MTP module applied recursively, Qwen3-Next style. It "
    "consumes the trunk hidden state at position i and the embedding of "
    "the token at i+k, and predicts the token at i+k+1. The LM head is "
    "shared with the trunk, which ties both computations to one output "
    "projection but does not guarantee close distributions. The module "
    "contains a transformer block whose attention was trained under a "
    "causal mask over the whole window: at inference it must be given "
    "the sequence, not a single position."
)

LAYER_PARAMETER_MAP = (
    ("attn_norm.weight", "input_layernorm.weight"),
    ("attn.wq.weight", "self_attn.q_proj.weight"),
    ("attn.wk.weight", "self_attn.k_proj.weight"),
    ("attn.wv.weight", "self_attn.v_proj.weight"),
    ("attn.wo.weight", "self_attn.o_proj.weight"),
    ("ffn_norm.weight", "post_attention_layernorm.weight"),
    ("ffn.w1.weight", "mlp.gate_proj.weight"),
    ("ffn.w3.weight", "mlp.up_proj.weight"),
    ("ffn.w2.weight", "mlp.down_proj.weight"),
)


class WispHFPackageError(ValueError):
    """The four-file Wisp HF package is missing or incompatible."""


class GreedyParityError(RuntimeError):
    """The MTP route did not reproduce the target greedy token stream."""


@dataclass(frozen=True)
class ReferenceDecodeResult:
    """Tokens and non-performance route telemetry from one reference decode."""

    token_ids: tuple[int, ...]
    generated_token_ids: tuple[int, ...]
    telemetry: dict[str, Any]

    def to_dict(self) -> dict[str, Any]:
        return {
            "token_ids": list(self.token_ids),
            "generated_token_ids": list(self.generated_token_ids),
            "telemetry": self.telemetry,
        }


def _duplicate_rejecting_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
    value: dict[str, Any] = {}
    for key, item in pairs:
        if key in value:
            raise WispHFPackageError(f"duplicate JSON key: {key!r}")
        value[key] = item
    return value


def _reject_json_constant(value: str) -> None:
    raise WispHFPackageError(f"non-finite JSON constant is forbidden: {value}")


def _checked_regular_file(package_dir: str, name: str) -> str:
    path = os.path.join(package_dir, name)
    try:
        info = os.lstat(path)
    except FileNotFoundError as exc:
        raise WispHFPackageError(f"required package file is missing: {name}") from exc
    if stat.S_ISLNK(info.st_mode):
        raise WispHFPackageError(f"package file must not be a symlink: {name}")
    if not stat.S_ISREG(info.st_mode):
        raise WispHFPackageError(f"package file is not regular: {name}")
    return path


def _file_sha256(path: str, label: str) -> str:
    """Hash one stable regular file without following a final symlink."""

    before = os.stat(path, follow_symlinks=False)
    flags = os.O_RDONLY
    if hasattr(os, "O_CLOEXEC"):
        flags |= os.O_CLOEXEC
    if hasattr(os, "O_NOFOLLOW"):
        flags |= os.O_NOFOLLOW
    try:
        descriptor = os.open(path, flags)
    except OSError as exc:
        raise WispHFPackageError(f"cannot hash {label}: {exc}") from exc
    digest = hashlib.sha256()
    try:
        with os.fdopen(descriptor, "rb") as handle:
            descriptor = -1
            while True:
                chunk = handle.read(1024 * 1024)
                if not chunk:
                    break
                digest.update(chunk)
            after = os.fstat(handle.fileno())
    finally:
        if descriptor >= 0:
            os.close(descriptor)
    identity_before = (
        before.st_dev,
        before.st_ino,
        before.st_size,
        before.st_mtime_ns,
    )
    identity_after = (
        after.st_dev,
        after.st_ino,
        after.st_size,
        after.st_mtime_ns,
    )
    if identity_before != identity_after:
        raise WispHFPackageError(f"{label} changed while it was being hashed")
    return digest.hexdigest()


def _load_verified_model_module(
    path: str,
    *,
    expected_sha256: str,
    expected_bytes: int,
) -> Any:
    """Execute the exact manifest-bound sibling source under a fresh name.

    The packaged runtime must not resolve ``wisp_mtp_model`` through
    ``sys.path`` or reuse a pre-existing ``sys.modules`` entry.  Read and hash
    the sibling ourselves, compile those exact bytes with their absolute path
    as the code origin, and expose the temporary module name only while its
    dataclasses are being defined.
    """

    absolute_path = os.path.abspath(path)
    if not _is_sha256(expected_sha256):
        raise WispHFPackageError(
            "export manifest wisp_mtp_model.py sha256 is invalid"
        )
    if (
        not isinstance(expected_bytes, int)
        or isinstance(expected_bytes, bool)
        or expected_bytes <= 0
    ):
        raise WispHFPackageError(
            "export manifest wisp_mtp_model.py byte count is invalid"
        )

    before = os.stat(absolute_path, follow_symlinks=False)
    if not stat.S_ISREG(before.st_mode):
        raise WispHFPackageError(
            "packaged model source is not a regular file"
        )
    if before.st_size != expected_bytes:
        raise WispHFPackageError(
            "packaged model source byte count does not match export manifest"
        )

    flags = os.O_RDONLY
    if hasattr(os, "O_CLOEXEC"):
        flags |= os.O_CLOEXEC
    if hasattr(os, "O_NOFOLLOW"):
        flags |= os.O_NOFOLLOW
    try:
        descriptor = os.open(absolute_path, flags)
    except OSError as exc:
        raise WispHFPackageError(
            f"cannot read packaged model source: {exc}"
        ) from exc

    try:
        with os.fdopen(descriptor, "rb") as handle:
            descriptor = -1
            opened = os.fstat(handle.fileno())
            source = handle.read()
            finished = os.fstat(handle.fileno())
    finally:
        if descriptor >= 0:
            os.close(descriptor)

    after = os.stat(absolute_path, follow_symlinks=False)
    identities = {
        (
            info.st_dev,
            info.st_ino,
            info.st_size,
            info.st_mtime_ns,
        )
        for info in (before, opened, finished, after)
    }
    if len(identities) != 1:
        raise WispHFPackageError(
            "packaged model source changed while it was being read"
        )
    if len(source) != expected_bytes:
        raise WispHFPackageError(
            "packaged model source read length does not match export manifest"
        )
    actual_sha256 = hashlib.sha256(source).hexdigest()
    if actual_sha256 != expected_sha256:
        raise WispHFPackageError(
            "packaged model source sha256 does not match export manifest"
        )

    unique_name = f"_wisp_mtp_model_{secrets.token_hex(16)}"
    if unique_name in sys.modules:
        raise WispHFPackageError(
            "fresh packaged model module name unexpectedly already exists"
        )
    spec = importlib.util.spec_from_file_location(unique_name, absolute_path)
    if (
        spec is None
        or spec.loader is None
        or os.path.abspath(str(spec.origin)) != absolute_path
    ):
        raise WispHFPackageError(
            "could not bind packaged model source to its absolute path"
        )
    module = importlib.util.module_from_spec(spec)
    try:
        code = compile(
            source,
            absolute_path,
            "exec",
            dont_inherit=True,
            optimize=0,
        )
        sys.modules[unique_name] = module
        exec(code, module.__dict__)
    except Exception as exc:
        raise WispHFPackageError(
            "packaged model source could not be executed"
        ) from exc
    finally:
        sys.modules.pop(unique_name, None)

    if (
        os.path.abspath(str(getattr(module, "__file__", ""))) != absolute_path
        or module.__spec__ is None
        or os.path.abspath(str(module.__spec__.origin)) != absolute_path
    ):
        raise WispHFPackageError(
            "packaged model module origin changed during execution"
        )
    if not isinstance(getattr(module, "ModelArgs", None), type):
        raise WispHFPackageError(
            "packaged model source does not export ModelArgs"
        )
    if not isinstance(getattr(module, "Wisp", None), type):
        raise WispHFPackageError("packaged model source does not export Wisp")
    if not callable(getattr(module, "causal_mask", None)):
        raise WispHFPackageError(
            "packaged model source does not export causal_mask"
        )
    return module


def _token_ids_sha256(token_ids: list[int] | tuple[int, ...]) -> str:
    digest = hashlib.sha256()
    digest.update(len(token_ids).to_bytes(8, "little", signed=False))
    for token in token_ids:
        digest.update(token.to_bytes(8, "little", signed=False))
    return digest.hexdigest()


def _read_strict_json(path: str, label: str) -> dict[str, Any]:
    before = os.stat(path, follow_symlinks=False)
    if before.st_size > 1024 * 1024:
        raise WispHFPackageError(f"{label} exceeds the 1 MiB metadata limit")
    flags = os.O_RDONLY
    if hasattr(os, "O_CLOEXEC"):
        flags |= os.O_CLOEXEC
    if hasattr(os, "O_NOFOLLOW"):
        flags |= os.O_NOFOLLOW
    try:
        descriptor = os.open(path, flags)
    except OSError as exc:
        raise WispHFPackageError(f"cannot open {label}: {exc}") from exc
    try:
        with os.fdopen(descriptor, "r", encoding="utf-8") as handle:
            descriptor = -1
            try:
                value = json.load(
                    handle,
                    object_pairs_hook=_duplicate_rejecting_object,
                    parse_constant=_reject_json_constant,
                )
            except (UnicodeDecodeError, json.JSONDecodeError) as exc:
                raise WispHFPackageError(f"{label} is not strict UTF-8 JSON") from exc
            after = os.fstat(handle.fileno())
    finally:
        if descriptor >= 0:
            os.close(descriptor)
    identity_before = (
        before.st_dev,
        before.st_ino,
        before.st_size,
        before.st_mtime_ns,
    )
    identity_after = (
        after.st_dev,
        after.st_ino,
        after.st_size,
        after.st_mtime_ns,
    )
    if identity_before != identity_after:
        raise WispHFPackageError(f"{label} changed while it was being read")
    if not isinstance(value, dict):
        raise WispHFPackageError(f"{label} must contain one JSON object")
    return value


def _is_sha256(value: Any) -> bool:
    return (
        isinstance(value, str)
        and len(value) == 64
        and all(character in "0123456789abcdef" for character in value)
    )


def _validate_repo_id(value: Any) -> str:
    if (
        not isinstance(value, str)
        or value != value.strip()
        or value.count("/") != 1
        or any(not part for part in value.split("/"))
        or any(character.isspace() for character in value)
    ):
        raise WispHFPackageError(
            "export_manifest.json repo_id must have form namespace/model"
        )
    return value


def _verify_export_manifest(package_dir: str) -> dict[str, Any]:
    """Locally verify the complete payload declared by export manifest v3."""

    manifest_path = _checked_regular_file(
        package_dir,
        "export_manifest.json",
    )
    manifest = _read_strict_json(
        manifest_path,
        "export_manifest.json",
    )
    actual_keys = frozenset(manifest)
    if actual_keys != EXPORT_MANIFEST_KEYS:
        missing = sorted(EXPORT_MANIFEST_KEYS - actual_keys)
        unexpected = sorted(actual_keys - EXPORT_MANIFEST_KEYS)
        raise WispHFPackageError(
            "export_manifest.json schema mismatch: "
            f"missing={missing}; unexpected={unexpected}"
        )
    if manifest["schema_version"] != 3:
        raise WispHFPackageError(
            "export_manifest.json schema_version must be 3"
        )
    repo_id = _validate_repo_id(manifest["repo_id"])
    release_complete = manifest["release_complete"]
    if not isinstance(release_complete, bool):
        raise WispHFPackageError(
            "export_manifest.json release_complete must be a boolean"
        )
    if (
        not isinstance(
            manifest["portable_evidence_bundle_required"],
            bool,
        )
        or manifest["portable_evidence_bundle_required"]
        is not release_complete
    ):
        raise WispHFPackageError(
            "portable evidence requirement must equal release completeness"
        )
    if not _is_sha256(manifest["model_card_template_sha256"]):
        raise WispHFPackageError(
            "model_card_template_sha256 is not a lowercase SHA-256"
        )

    evaluation_sources = manifest["evaluation_sources"]
    if release_complete:
        if (
            not isinstance(evaluation_sources, dict)
            or frozenset(evaluation_sources) != EVALUATION_SOURCE_KEYS
        ):
            raise WispHFPackageError(
                "complete export has an invalid evaluation_sources set"
            )
        for name in sorted(EVALUATION_SOURCE_KEYS):
            evidence = evaluation_sources[name]
            if (
                not isinstance(evidence, dict)
                or frozenset(evidence) != {"sha256"}
                or not _is_sha256(evidence["sha256"])
            ):
                raise WispHFPackageError(
                    f"evaluation source {name!r} is not hash-bound"
                )
    elif evaluation_sources is not None:
        raise WispHFPackageError(
            "development export must not declare evaluation_sources"
        )

    source_checkpoint = manifest["source_checkpoint"]
    if (
        not isinstance(source_checkpoint, dict)
        or frozenset(source_checkpoint)
        != {
            "step",
            "meta_sha256",
            "master_sha256",
            "optimizer_sha256",
        }
    ):
        raise WispHFPackageError(
            "export_manifest.json source_checkpoint schema is invalid"
        )
    step = source_checkpoint["step"]
    if not isinstance(step, int) or isinstance(step, bool) or step < 1:
        raise WispHFPackageError(
            "export_manifest.json source checkpoint step is invalid"
        )
    for name in ("meta_sha256", "master_sha256", "optimizer_sha256"):
        if not _is_sha256(source_checkpoint[name]):
            raise WispHFPackageError(
                f"source_checkpoint.{name} is not a lowercase SHA-256"
            )

    for name in (
        "trunk_parameters",
        "mtp_parameters_excluding_shared_embedding_and_head",
    ):
        value = manifest[name]
        if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
            raise WispHFPackageError(
                f"export_manifest.json {name} must be a positive integer"
            )

    declared_files = manifest["files"]
    if (
        not isinstance(declared_files, dict)
        or tuple(sorted(declared_files)) != EXPECTED_RELEASE_PAYLOAD_FILES
    ):
        actual = (
            sorted(declared_files)
            if isinstance(declared_files, dict)
            else type(declared_files).__name__
        )
        raise WispHFPackageError(
            "export manifest payload differs from the packaged runtime "
            f"contract: {actual}"
        )
    for name in EXPECTED_RELEASE_PAYLOAD_FILES:
        evidence = declared_files[name]
        if (
            not isinstance(evidence, dict)
            or frozenset(evidence) != {"bytes", "sha256"}
            or not isinstance(evidence["bytes"], int)
            or isinstance(evidence["bytes"], bool)
            or evidence["bytes"] < 0
            or not _is_sha256(evidence["sha256"])
        ):
            raise WispHFPackageError(
                f"manifest evidence for {name!r} is malformed"
            )
        path = _checked_regular_file(package_dir, name)
        if (
            os.path.getsize(path) != evidence["bytes"]
            or _file_sha256(path, name) != evidence["sha256"]
        ):
            raise WispHFPackageError(
                f"release artifact {name!r} does not match export manifest"
            )

    readme_path = os.path.join(package_dir, "README.md")
    try:
        with open(readme_path, encoding="utf-8") as handle:
            readme = handle.read()
    except (OSError, UnicodeDecodeError) as exc:
        raise WispHFPackageError("README.md is not readable UTF-8") from exc
    if (
        "{{REPO_ID}}" in readme
        or "{{FINAL_EVALUATION}}" in readme
        or repo_id not in readme
    ):
        raise WispHFPackageError(
            "README.md does not match export manifest repo_id"
        )
    return manifest


def _require_exact_keys(
    value: dict[str, Any],
    expected: frozenset[str],
    label: str,
) -> None:
    actual = frozenset(value)
    missing = sorted(expected - actual)
    unexpected = sorted(actual - expected)
    if missing or unexpected:
        details = []
        if missing:
            details.append(f"missing={missing}")
        if unexpected:
            details.append(f"unexpected={unexpected}")
        raise WispHFPackageError(f"{label} schema mismatch: {'; '.join(details)}")


def _positive_int(value: Any, label: str) -> int:
    if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
        raise WispHFPackageError(f"{label} must be a positive integer")
    return value


def _nonnegative_int(value: Any, label: str) -> int:
    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
        raise WispHFPackageError(f"{label} must be a non-negative integer")
    return value


def _positive_number(value: Any, label: str) -> float:
    if (
        not isinstance(value, (int, float))
        or isinstance(value, bool)
        or not math.isfinite(value)
        or value <= 0
    ):
        raise WispHFPackageError(f"{label} must be a positive finite number")
    return float(value)


def _require_literal(value: Any, expected: Any, label: str) -> None:
    if type(value) is not type(expected) or value != expected:
        raise WispHFPackageError(f"{label} must be exactly {expected!r}")


def _model_args_from_configs(
    model_config: dict[str, Any],
    mtp_config: dict[str, Any],
    model_args_type: type[Any],
) -> Any:
    _require_exact_keys(model_config, MODEL_CONFIG_KEYS, "config.json")
    _require_exact_keys(mtp_config, MTP_CONFIG_KEYS, "mtp_config.json")

    _require_literal(
        model_config["architectures"],
        ["LlamaForCausalLM"],
        "config.json architectures",
    )
    _require_literal(model_config["model_type"], "llama", "config.json model_type")
    _require_literal(model_config["hidden_act"], "silu", "config.json hidden_act")
    _require_literal(
        model_config["attention_bias"],
        False,
        "config.json attention_bias",
    )
    _require_literal(model_config["mlp_bias"], False, "config.json mlp_bias")
    _require_literal(
        model_config["torch_dtype"],
        "bfloat16",
        "config.json torch_dtype",
    )
    _require_literal(model_config["bos_token_id"], None, "config.json bos_token_id")
    _require_literal(model_config["eos_token_id"], 0, "config.json eos_token_id")
    _require_literal(model_config["pad_token_id"], 1, "config.json pad_token_id")
    if not isinstance(model_config["tie_word_embeddings"], bool):
        raise WispHFPackageError(
            "config.json tie_word_embeddings must be a boolean"
        )

    dim = _positive_int(model_config["hidden_size"], "hidden_size")
    ffn_hidden = _positive_int(
        model_config["intermediate_size"],
        "intermediate_size",
    )
    n_layers = _positive_int(
        model_config["num_hidden_layers"],
        "num_hidden_layers",
    )
    n_heads = _positive_int(
        model_config["num_attention_heads"],
        "num_attention_heads",
    )
    n_kv_heads = _positive_int(
        model_config["num_key_value_heads"],
        "num_key_value_heads",
    )
    head_dim = _positive_int(model_config["head_dim"], "head_dim")
    max_seq_len = _positive_int(
        model_config["max_position_embeddings"],
        "max_position_embeddings",
    )
    vocab_size = _positive_int(model_config["vocab_size"], "vocab_size")
    norm_eps = _positive_number(model_config["rms_norm_eps"], "rms_norm_eps")
    rope_theta = _positive_number(model_config["rope_theta"], "rope_theta")
    mtp_layers = _positive_int(mtp_config["mtp_layers"], "mtp_layers")
    mtp_depth = _positive_int(
        mtp_config["mtp_depth_trained"],
        "mtp_depth_trained",
    )

    _require_literal(
        mtp_config["shared_lm_head"],
        True,
        "mtp_config.json shared_lm_head",
    )
    _require_literal(
        mtp_config["recursive"],
        True,
        "mtp_config.json recursive",
    )
    _require_literal(
        mtp_config["schema_version"],
        MTP_SCHEMA_VERSION,
        "mtp_config.json schema_version",
    )
    _require_literal(
        mtp_config["architecture"],
        MTP_ARCHITECTURE,
        "mtp_config.json architecture",
    )
    _require_literal(
        mtp_config["hidden_state_stage"],
        MTP_HIDDEN_STATE_STAGE,
        "mtp_config.json hidden_state_stage",
    )
    _require_literal(
        mtp_config["requires_full_sequence_attention"],
        MTP_REQUIRES_FULL_SEQUENCE_ATTENTION,
        "mtp_config.json requires_full_sequence_attention",
    )
    _require_literal(
        mtp_config["tensor_prefix"],
        MTP_TENSOR_PREFIX,
        "mtp_config.json tensor_prefix",
    )
    # Schema v1 is machine-readable, but retain the exact note as part of the
    # versioned contract so a producer cannot change the detailed position
    # semantics without a schema bump.
    _require_literal(
        mtp_config["note"],
        EXPECTED_MTP_NOTE,
        "mtp_config.json note",
    )
    _nonnegative_int(mtp_config["trained_steps"], "trained_steps")

    if vocab_size < 2:
        raise WispHFPackageError("vocab_size must include token IDs 0 and 1")
    if max_seq_len < 2:
        raise WispHFPackageError("max_position_embeddings must be at least 2")
    if dim % n_heads:
        raise WispHFPackageError(
            "hidden_size must be divisible by num_attention_heads"
        )
    if n_heads % n_kv_heads:
        raise WispHFPackageError(
            "num_attention_heads must be divisible by num_key_value_heads"
        )
    if head_dim != dim // n_heads:
        raise WispHFPackageError(
            "head_dim does not equal hidden_size / num_attention_heads"
        )
    if head_dim % 2:
        raise WispHFPackageError("head_dim must be even for Wisp RoPE")

    return model_args_type(
        vocab_size=vocab_size,
        dim=dim,
        n_layers=n_layers,
        n_heads=n_heads,
        n_kv_heads=n_kv_heads,
        ffn_hidden=ffn_hidden,
        max_seq_len=max_seq_len,
        rope_theta=rope_theta,
        norm_eps=norm_eps,
        tie_embeddings=model_config["tie_word_embeddings"],
        mtp_layers=mtp_layers,
        mtp_depth=mtp_depth,
        ce_chunk=0,
    )


def _block_shapes(prefix: str, args: Any) -> dict[str, tuple[int, ...]]:
    query_width = args.n_heads * args.head_dim
    kv_width = args.n_kv_heads * args.head_dim
    return {
        f"{prefix}.attn_norm.weight": (args.dim,),
        f"{prefix}.attn.wq.weight": (query_width, args.dim),
        f"{prefix}.attn.wk.weight": (kv_width, args.dim),
        f"{prefix}.attn.wv.weight": (kv_width, args.dim),
        f"{prefix}.attn.wo.weight": (args.dim, query_width),
        f"{prefix}.ffn_norm.weight": (args.dim,),
        f"{prefix}.ffn.w1.weight": (args.ffn_hidden, args.dim),
        f"{prefix}.ffn.w3.weight": (args.ffn_hidden, args.dim),
        f"{prefix}.ffn.w2.weight": (args.dim, args.ffn_hidden),
    }


def _expected_internal_trunk_shapes(
    args: Any,
) -> dict[str, tuple[int, ...]]:
    expected = {
        "tok_emb.weight": (args.vocab_size, args.dim),
        "norm.weight": (args.dim,),
    }
    if not args.tie_embeddings:
        expected["lm_head.weight"] = (args.vocab_size, args.dim)
    for layer in range(args.n_layers):
        expected.update(_block_shapes(f"blocks.{layer}", args))
    return expected


def _expected_hf_trunk_shapes(
    args: Any,
) -> dict[str, tuple[int, ...]]:
    internal = _expected_internal_trunk_shapes(args)
    expected = {
        "model.embed_tokens.weight": internal["tok_emb.weight"],
        "model.norm.weight": internal["norm.weight"],
    }
    if not args.tie_embeddings:
        expected["lm_head.weight"] = internal["lm_head.weight"]
    for layer in range(args.n_layers):
        for ours, theirs in LAYER_PARAMETER_MAP:
            expected[f"model.layers.{layer}.{theirs}"] = internal[
                f"blocks.{layer}.{ours}"
            ]
    return expected


def _expected_mtp_shapes(args: Any) -> dict[str, tuple[int, ...]]:
    expected = {
        "mtp.h_norm.weight": (args.dim,),
        "mtp.e_norm.weight": (args.dim,),
        "mtp.proj.weight": (args.dim, 2 * args.dim),
    }
    for layer in range(args.mtp_layers):
        expected.update(_block_shapes(f"mtp.blocks.{layer}", args))
    return expected


def _load_and_validate_tensors(
    path: str,
    expected: dict[str, tuple[int, ...]],
    label: str,
) -> dict[str, mx.array]:
    before = os.stat(path, follow_symlinks=False)
    try:
        tensors = mx.load(path)
    except Exception as exc:
        raise WispHFPackageError(f"{label} is not a readable safetensors file") from exc
    if not isinstance(tensors, dict) or not all(
        isinstance(key, str) and isinstance(value, mx.array)
        for key, value in tensors.items()
    ):
        raise WispHFPackageError(f"{label} did not contain a tensor dictionary")

    actual_names = set(tensors)
    expected_names = set(expected)
    missing = sorted(expected_names - actual_names)
    unexpected = sorted(actual_names - expected_names)
    if missing or unexpected:
        details = []
        if missing:
            details.append(f"missing={missing}")
        if unexpected:
            details.append(f"unexpected={unexpected}")
        raise WispHFPackageError(
            f"{label} tensor set mismatch: {'; '.join(details)}"
        )

    for name in sorted(expected):
        tensor = tensors[name]
        actual_shape = tuple(int(value) for value in tensor.shape)
        if actual_shape != expected[name]:
            raise WispHFPackageError(
                f"{label} tensor {name!r} has shape {actual_shape}, "
                f"expected {expected[name]}"
            )
        if tensor.dtype != mx.bfloat16:
            raise WispHFPackageError(
                f"{label} tensor {name!r} has dtype {tensor.dtype}, "
                "expected mlx.core.bfloat16"
            )

    finite_checks = [
        (name, mx.all(mx.isfinite(tensors[name])))
        for name in sorted(expected)
    ]
    mx.eval(*[check for _, check in finite_checks])
    non_finite = [
        name for name, check in finite_checks if not bool(check.item())
    ]
    if non_finite:
        raise WispHFPackageError(
            f"{label} contains non-finite tensors: {non_finite}"
        )

    after = os.stat(path, follow_symlinks=False)
    identity_before = (
        before.st_dev,
        before.st_ino,
        before.st_size,
        before.st_mtime_ns,
    )
    identity_after = (
        after.st_dev,
        after.st_ino,
        after.st_size,
        after.st_mtime_ns,
    )
    if identity_before != identity_after:
        raise WispHFPackageError(f"{label} changed while it was being loaded")
    return tensors


def _hf_trunk_to_internal(
    tensors: dict[str, mx.array],
    args: Any,
) -> dict[str, mx.array]:
    internal = {
        "tok_emb.weight": tensors["model.embed_tokens.weight"],
        "norm.weight": tensors["model.norm.weight"],
    }
    if not args.tie_embeddings:
        internal["lm_head.weight"] = tensors["lm_head.weight"]
    for layer in range(args.n_layers):
        for ours, theirs in LAYER_PARAMETER_MAP:
            internal[f"blocks.{layer}.{ours}"] = tensors[
                f"model.layers.{layer}.{theirs}"
            ]
    return internal


class WispMTPReferenceRuntime:
    """Loaded Wisp HF package plus greedy correctness routes."""

    def __init__(
        self,
        package_dir: str,
        model: Any,
        args: Any,
        model_module: Any,
        model_config: dict[str, Any],
        mtp_config: dict[str, Any],
        package_file_sha256: dict[str, str],
        export_manifest: dict[str, Any],
    ) -> None:
        self.package_dir = package_dir
        self.model = model
        self.args = args
        self._model_module = model_module
        self._causal_mask = model_module.causal_mask
        self.model_config = model_config
        self.mtp_config = mtp_config
        self.package_file_sha256 = package_file_sha256
        self.export_manifest = export_manifest

    @classmethod
    def load(cls, package_dir: str) -> "WispMTPReferenceRuntime":
        package_dir = os.path.abspath(package_dir)
        try:
            package_info = os.lstat(package_dir)
        except FileNotFoundError as exc:
            raise WispHFPackageError(
                f"package directory does not exist: {package_dir}"
            ) from exc
        if stat.S_ISLNK(package_info.st_mode):
            raise WispHFPackageError("package directory must not be a symlink")
        if not stat.S_ISDIR(package_info.st_mode):
            raise WispHFPackageError("package path must be a directory")

        paths = {
            name: _checked_regular_file(package_dir, name)
            for name in REQUIRED_PACKAGE_FILES
        }
        digests_before = {
            name: _file_sha256(path, name)
            for name, path in paths.items()
        }
        strict_manifest = _verify_export_manifest(package_dir)
        model_source_evidence = strict_manifest["files"][
            "wisp_mtp_model.py"
        ]
        model_module = _load_verified_model_module(
            _checked_regular_file(package_dir, "wisp_mtp_model.py"),
            expected_sha256=model_source_evidence["sha256"],
            expected_bytes=model_source_evidence["bytes"],
        )
        model_config = _read_strict_json(paths["config.json"], "config.json")
        mtp_config = _read_strict_json(
            paths["mtp_config.json"],
            "mtp_config.json",
        )
        args = _model_args_from_configs(
            model_config,
            mtp_config,
            model_module.ModelArgs,
        )
        source_checkpoint = strict_manifest.get("source_checkpoint")
        if (
            not isinstance(source_checkpoint, dict)
            or source_checkpoint.get("step") != mtp_config["trained_steps"]
        ):
            raise WispHFPackageError(
                "mtp_config.json trained_steps does not match "
                "export_manifest.json source_checkpoint.step"
            )
        trunk = _load_and_validate_tensors(
            paths["model.safetensors"],
            _expected_hf_trunk_shapes(args),
            "model.safetensors",
        )
        sidecar = _load_and_validate_tensors(
            paths["mtp.safetensors"],
            _expected_mtp_shapes(args),
            "mtp.safetensors",
        )
        trunk_parameters = sum(int(value.size) for value in trunk.values())
        mtp_parameters = sum(int(value.size) for value in sidecar.values())
        if strict_manifest.get("trunk_parameters") != trunk_parameters:
            raise WispHFPackageError(
                "model.safetensors parameter count does not match "
                "export_manifest.json"
            )
        if (
            strict_manifest.get(
                "mtp_parameters_excluding_shared_embedding_and_head"
            )
            != mtp_parameters
        ):
            raise WispHFPackageError(
                "mtp.safetensors parameter count does not match "
                "export_manifest.json"
            )

        internal = _hf_trunk_to_internal(trunk, args)
        internal.update(sidecar)
        expected_internal = _expected_internal_trunk_shapes(args)
        expected_internal.update(_expected_mtp_shapes(args))
        if set(internal) != set(expected_internal):
            raise WispHFPackageError(
                "internal parameter mapping did not cover the Wisp model exactly"
            )

        model = model_module.Wisp(args)
        model.update(tree_unflatten(list(sorted(internal.items()))))
        model.eval()
        mx.eval(model.parameters())
        loaded_parameters = dict(tree_flatten(model.parameters()))
        if set(loaded_parameters) != set(expected_internal):
            raise WispHFPackageError(
                "loaded Wisp parameter tree differs from the package contract"
            )
        for name, expected_shape in expected_internal.items():
            tensor = loaded_parameters[name]
            if tuple(tensor.shape) != expected_shape or tensor.dtype != mx.bfloat16:
                raise WispHFPackageError(
                    f"loaded parameter {name!r} changed shape or dtype"
                )
        digests_after = {
            name: _file_sha256(path, name)
            for name, path in paths.items()
        }
        if digests_before != digests_after:
            raise WispHFPackageError(
                "package files changed while the runtime was loading"
            )
        verified_after = _verify_export_manifest(package_dir)
        if verified_after != strict_manifest:
            raise WispHFPackageError(
                "export manifest payload changed while the runtime was loading"
            )
        payload_hashes = {
            name: evidence["sha256"]
            for name, evidence in strict_manifest["files"].items()
        }
        payload_hashes["export_manifest.json"] = digests_after[
            "export_manifest.json"
        ]
        return cls(
            package_dir,
            model,
            args,
            model_module,
            model_config,
            mtp_config,
            payload_hashes,
            strict_manifest,
        )

    def package_summary(self) -> dict[str, Any]:
        return {
            "format": "wisp_hf_split_mtp_v1",
            "mtp_schema_version": MTP_SCHEMA_VERSION,
            "mtp_architecture": MTP_ARCHITECTURE,
            "hidden_state_stage": MTP_HIDDEN_STATE_STAGE,
            "requires_full_sequence_attention": (
                MTP_REQUIRES_FULL_SEQUENCE_ATTENTION
            ),
            "tensor_prefix": MTP_TENSOR_PREFIX,
            "shared_lm_head": True,
            "recursive": True,
            "metadata_semantics_gate": "mtp_schema_v1_and_exact_note",
            "lineage_scope": "export_manifest_v3_bound_payload",
            "package_lineage_verified": True,
            "provenance_attested": False,
            "package_file_sha256": dict(self.package_file_sha256),
            "export_manifest_sha256": self.package_file_sha256[
                "export_manifest.json"
            ],
            "repo_id": self.export_manifest["repo_id"],
            "source_checkpoint": dict(
                self.export_manifest["source_checkpoint"]
            ),
            "model_type": "llama",
            "vocab_size": self.args.vocab_size,
            "hidden_size": self.args.dim,
            "num_hidden_layers": self.args.n_layers,
            "num_attention_heads": self.args.n_heads,
            "num_key_value_heads": self.args.n_kv_heads,
            "mtp_layers": self.args.mtp_layers,
            "mtp_depth_trained": self.args.mtp_depth,
            "trained_steps": self.mtp_config["trained_steps"],
            "runtime_kind": "correctness_reference",
            "production_latency_claim": False,
        }

    def _validate_prompt(self, prompt_ids: list[int] | tuple[int, ...]) -> list[int]:
        if not isinstance(prompt_ids, (list, tuple)) or not prompt_ids:
            raise ValueError("prompt_ids must be a non-empty list or tuple")
        validated = []
        for index, token in enumerate(prompt_ids):
            if not isinstance(token, int) or isinstance(token, bool):
                raise ValueError(f"prompt_ids[{index}] is not an integer")
            if token < 0 or token >= self.args.vocab_size:
                raise ValueError(
                    f"prompt_ids[{index}]={token} is outside the vocabulary"
                )
            validated.append(token)
        if len(validated) > self.args.max_seq_len:
            raise ValueError("prompt exceeds max_position_embeddings")
        return validated

    def _validate_decode(
        self,
        prompt_ids: list[int] | tuple[int, ...],
        max_new_tokens: int,
    ) -> list[int]:
        tokens = self._validate_prompt(prompt_ids)
        if (
            not isinstance(max_new_tokens, int)
            or isinstance(max_new_tokens, bool)
            or max_new_tokens < 0
        ):
            raise ValueError("max_new_tokens must be a non-negative integer")
        if len(tokens) + max_new_tokens > self.args.max_seq_len:
            raise ValueError(
                "prompt plus max_new_tokens exceeds max_position_embeddings"
            )
        return tokens

    def _validate_depth(self, depth: int | None) -> int:
        if depth is None:
            return self.args.mtp_depth
        if not isinstance(depth, int) or isinstance(depth, bool) or depth <= 0:
            raise ValueError("depth must be a positive integer")
        if depth > self.args.mtp_depth:
            raise ValueError(
                "depth exceeds mtp_depth_trained; this correctness runtime "
                "does not make untrained-depth claims"
            )
        return depth

    def target_logits(
        self,
        token_ids: list[int] | tuple[int, ...],
    ) -> mx.array:
        """Return full target logits for a validated sequence."""

        tokens = self._validate_prompt(token_ids)
        sequence = mx.array([tokens], dtype=mx.int32)
        mask = (
            self._causal_mask(len(tokens), self.model.norm.weight.dtype)
            if len(tokens) > 1
            else None
        )
        logits, _, _ = self.model(sequence, mask)
        mx.eval(logits)
        return logits

    @staticmethod
    def _greedy_token(logits: mx.array) -> int:
        return int(mx.argmax(logits).item())

    def _telemetry_base(self, route: str) -> dict[str, Any]:
        return {
            "runtime_kind": "correctness_reference",
            "production_latency_claim": False,
            "latency_measurement": None,
            "throughput_measurement": None,
            "route": route,
            "greedy": True,
            "token_parity_rule": "exact_token_ids",
            "near_tie_tolerance": False,
            "lineage_scope": "export_manifest_v3_bound_payload",
            "package_lineage_verified": True,
            "provenance_attested": False,
            "package_file_sha256": dict(self.package_file_sha256),
            "export_manifest_sha256": self.package_file_sha256[
                "export_manifest.json"
            ],
            "source_checkpoint": dict(
                self.export_manifest["source_checkpoint"]
            ),
        }

    def decode_ar(
        self,
        prompt_ids: list[int] | tuple[int, ...],
        max_new_tokens: int,
    ) -> ReferenceDecodeResult:
        """Greedy target AR using the repository's ordinary KV-cache route."""

        tokens = self._validate_decode(prompt_ids, max_new_tokens)
        prompt_length = len(tokens)
        caches = None
        fed = mx.array([tokens], dtype=mx.int32)
        prefill_forwards = 0
        decode_forwards = 0
        for generated_index in range(max_new_tokens):
            length = int(fed.shape[1])
            mask = (
                self._causal_mask(length, self.model.norm.weight.dtype)
                if caches is None and length > 1
                else None
            )
            logits, _, caches = self.model(fed, mask, caches)
            mx.eval(logits, caches)
            token = self._greedy_token(logits[0, -1])
            tokens.append(token)
            fed = mx.array([[token]], dtype=mx.int32)
            if generated_index == 0:
                prefill_forwards += 1
            else:
                decode_forwards += 1

        telemetry = self._telemetry_base("target_ar_kv_cache")
        telemetry.update(
            {
                "uses_target_model": True,
                "uses_mtp_module": False,
                "mtp_route_selected": False,
                "mtp_module_executed": False,
                "uses_target_kv_cache": True,
                "rollback_capable_target_cache": False,
                "prompt_tokens": prompt_length,
                "requested_new_tokens": max_new_tokens,
                "generated_tokens": max_new_tokens,
                "target_prefill_forwards": prefill_forwards,
                "target_decode_forwards": decode_forwards,
                "target_verification_forwards": 0,
                "target_forwards": prefill_forwards + decode_forwards,
                "mtp_recursions": 0,
                "drafts_issued": 0,
                "drafts_accepted": 0,
                "drafts_rejected": 0,
                "prompt_token_ids_sha256": _token_ids_sha256(
                    tuple(tokens[:prompt_length])
                ),
                "output_token_ids_sha256": _token_ids_sha256(tuple(tokens)),
                "generated_token_ids_sha256": _token_ids_sha256(
                    tuple(tokens[prompt_length:])
                ),
            }
        )
        return ReferenceDecodeResult(
            token_ids=tuple(tokens),
            generated_token_ids=tuple(tokens[prompt_length:]),
            telemetry=telemetry,
        )

    def decode_mtp(
        self,
        prompt_ids: list[int] | tuple[int, ...],
        max_new_tokens: int,
        depth: int | None = None,
    ) -> ReferenceDecodeResult:
        """Greedy self-speculation with full-prefix target verification.

        Each cycle emits one target token, drafts up to ``depth`` more tokens
        through the recursively shared MTP module, and verifies all drafts with
        the target trunk.  A mismatching draft is replaced with the target
        argmax.  The route is intentionally not optimized for wall-clock speed.
        """

        tokens = self._validate_decode(prompt_ids, max_new_tokens)
        draft_depth = self._validate_depth(depth)
        prompt_length = len(tokens)
        produced = 0
        prefix_hidden = None
        next_target_logits = None

        cycles = 0
        target_prefix_forwards = 0
        target_recovery_forwards = 0
        target_verification_forwards = 0
        mtp_recursions = 0
        drafts_issued = 0
        drafts_accepted = 0
        drafts_rejected = 0
        unused_drafts_after_rejection = 0
        fully_accepted_verifications = 0
        draft_trials_per_depth = [0 for _ in range(draft_depth)]
        draft_accepts_per_depth = [0 for _ in range(draft_depth)]

        while produced < max_new_tokens:
            cycles += 1
            if prefix_hidden is None:
                sequence = mx.array([tokens], dtype=mx.int32)
                mask = (
                    self._causal_mask(
                        len(tokens),
                        self.model.norm.weight.dtype,
                    )
                    if len(tokens) > 1
                    else None
                )
                prefix_hidden, _ = self.model.trunk(sequence, mask)
                next_logits = self.model.head(prefix_hidden[:, -1:, :])
                mx.eval(prefix_hidden, next_logits)
                next_target_logits = next_logits[0, -1]
                target_prefix_forwards += 1
                if cycles > 1:
                    target_recovery_forwards += 1

            bonus = self._greedy_token(next_target_logits)
            tokens.append(bonus)
            produced += 1
            if produced >= max_new_tokens:
                break

            issue_count = min(draft_depth, max_new_tokens - produced)
            prefix_length = int(prefix_hidden.shape[1])
            mtp_mask = self._causal_mask(
                prefix_length,
                self.model.norm.weight.dtype,
            )
            current_hidden = prefix_hidden
            conditioning_window = tokens[1 : prefix_length + 1]
            drafts = []
            for _ in range(issue_count):
                token_embeddings = self.model.tok_emb(
                    mx.array([conditioning_window], dtype=mx.int32)
                )
                current_hidden, _ = self.model.mtp(
                    current_hidden,
                    token_embeddings,
                    mtp_mask,
                )
                draft_logits = self.model.head(current_hidden[:, -1:, :])
                mx.eval(current_hidden, draft_logits)
                draft = self._greedy_token(draft_logits[0, -1])
                drafts.append(draft)
                conditioning_window = conditioning_window[1:] + [draft]
                mtp_recursions += 1
            drafts_issued += len(drafts)

            candidate = tokens + drafts
            sequence = mx.array([candidate], dtype=mx.int32)
            verify_mask = self._causal_mask(
                len(candidate),
                self.model.norm.weight.dtype,
            )
            verified_hidden, _ = self.model.trunk(sequence, verify_mask)
            base = len(tokens) - 1
            verify_logits = self.model.head(
                verified_hidden[:, base : base + len(drafts) + 1, :]
            )
            mx.eval(verified_hidden, verify_logits)
            target_verification_forwards += 1

            accepted_this_verification = 0
            rejected = False
            for draft_index, draft in enumerate(drafts):
                draft_trials_per_depth[draft_index] += 1
                target_token = self._greedy_token(verify_logits[0, draft_index])
                if draft == target_token:
                    tokens.append(draft)
                    produced += 1
                    drafts_accepted += 1
                    draft_accepts_per_depth[draft_index] += 1
                    accepted_this_verification += 1
                else:
                    tokens.append(target_token)
                    produced += 1
                    drafts_rejected += 1
                    unused_drafts_after_rejection += (
                        len(drafts) - draft_index - 1
                    )
                    rejected = True
                if rejected or produced >= max_new_tokens:
                    break

            fully_accepted = accepted_this_verification == len(drafts)
            if fully_accepted:
                fully_accepted_verifications += 1
                if produced < max_new_tokens:
                    prefix_hidden = verified_hidden
                    next_target_logits = verify_logits[0, -1]
            elif produced < max_new_tokens:
                prefix_hidden = None
                next_target_logits = None

        drafts_trialled = sum(draft_trials_per_depth)
        if drafts_trialled != drafts_accepted + drafts_rejected:
            raise RuntimeError("internal MTP trial accounting invariant failed")
        if drafts_issued != drafts_trialled + unused_drafts_after_rejection:
            raise RuntimeError("internal MTP issuance accounting invariant failed")

        telemetry = self._telemetry_base("wisp_mtp_full_prefix_reference")
        telemetry.update(
            {
                "uses_target_model": True,
                "uses_mtp_module": mtp_recursions > 0,
                "mtp_route_selected": True,
                "mtp_module_executed": mtp_recursions > 0,
                "uses_target_kv_cache": False,
                "rollback_capable_target_cache": False,
                "corrected_prefix_recompute_supported": True,
                "corrected_prefix_recomputes": target_recovery_forwards,
                "prompt_tokens": prompt_length,
                "requested_new_tokens": max_new_tokens,
                "generated_tokens": produced,
                "draft_depth": draft_depth,
                "mtp_depth_trained": self.args.mtp_depth,
                "cycles": cycles,
                "target_prefix_forwards": target_prefix_forwards,
                "target_recovery_forwards": target_recovery_forwards,
                "target_decode_forwards": 0,
                "target_verification_forwards": target_verification_forwards,
                "target_forwards": (
                    target_prefix_forwards + target_verification_forwards
                ),
                "mtp_recursions": mtp_recursions,
                "drafts_issued": drafts_issued,
                "drafts_trialled": drafts_trialled,
                "drafts_accepted": drafts_accepted,
                "drafts_rejected": drafts_rejected,
                "unused_drafts_after_rejection": (
                    unused_drafts_after_rejection
                ),
                "fully_accepted_verifications": (
                    fully_accepted_verifications
                ),
                "draft_trials_per_depth": draft_trials_per_depth,
                "draft_accepts_per_depth": draft_accepts_per_depth,
                "prompt_token_ids_sha256": _token_ids_sha256(
                    tuple(tokens[:prompt_length])
                ),
                "output_token_ids_sha256": _token_ids_sha256(tuple(tokens)),
                "generated_token_ids_sha256": _token_ids_sha256(
                    tuple(tokens[prompt_length:])
                ),
            }
        )
        return ReferenceDecodeResult(
            token_ids=tuple(tokens),
            generated_token_ids=tuple(tokens[prompt_length:]),
            telemetry=telemetry,
        )

    def verify_greedy_parity(
        self,
        prompt_ids: list[int] | tuple[int, ...],
        max_new_tokens: int,
        depth: int | None = None,
    ) -> dict[str, Any]:
        """Run both routes and fail unless MTP was exercised and tokens match."""

        if (
            not isinstance(max_new_tokens, int)
            or isinstance(max_new_tokens, bool)
            or max_new_tokens < 2
        ):
            raise ValueError(
                "greedy parity needs at least two new tokens to exercise MTP"
            )
        ar = self.decode_ar(prompt_ids, max_new_tokens)
        mtp = self.decode_mtp(prompt_ids, max_new_tokens, depth)
        if mtp.telemetry["mtp_recursions"] <= 0:
            raise GreedyParityError(
                "MTP parity route did not execute an MTP recursion"
            )
        if ar.token_ids != mtp.token_ids:
            overlap = min(len(ar.token_ids), len(mtp.token_ids))
            mismatch = next(
                (
                    index for index, (target_token, mtp_token) in enumerate(
                        zip(
                            ar.token_ids[:overlap],
                            mtp.token_ids[:overlap],
                        )
                    ) if target_token != mtp_token
                ),
                overlap,
            )
            ar_value = (
                ar.token_ids[mismatch]
                if mismatch < len(ar.token_ids)
                else None
            )
            mtp_value = (
                mtp.token_ids[mismatch] if mismatch < len(mtp.token_ids) else None
            )
            raise GreedyParityError(
                "exact greedy parity failed at absolute token index "
                f"{mismatch}: AR={ar_value}, MTP={mtp_value}; "
                f"lengths AR={len(ar.token_ids)}, MTP={len(mtp.token_ids)}"
            )
        ar.telemetry["exact_greedy_parity_verified"] = True
        mtp.telemetry["exact_greedy_parity_verified"] = True
        return {
            "exact_greedy_parity": True,
            "token_ids": list(ar.token_ids),
            "generated_token_ids": list(ar.generated_token_ids),
            "ar_route": ar.telemetry,
            "mtp_route": mtp.telemetry,
        }


def _parse_prompt_ids(value: str) -> list[int]:
    text = value.strip()
    if not text:
        raise argparse.ArgumentTypeError("prompt IDs must not be empty")
    try:
        if text.startswith("["):
            parsed = json.loads(
                text,
                parse_constant=lambda item: (_ for _ in ()).throw(
                    ValueError(f"non-finite constant {item}")
                ),
            )
        else:
            parsed = [int(item.strip()) for item in text.split(",")]
    except (ValueError, json.JSONDecodeError) as exc:
        raise argparse.ArgumentTypeError(
            "prompt IDs must be a JSON array or comma-separated integers"
        ) from exc
    if not isinstance(parsed, list):
        raise argparse.ArgumentTypeError("prompt IDs must form a list")
    if not all(isinstance(item, int) and not isinstance(item, bool) for item in parsed):
        raise argparse.ArgumentTypeError("every prompt ID must be an integer")
    return parsed


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Load a manifest-bound Wisp HF split export and run a greedy "
            "correctness route. No production latency is measured."
        )
    )
    parser.add_argument("--package", required=True)
    parser.add_argument("--prompt-ids", required=True, type=_parse_prompt_ids)
    parser.add_argument(
        "--mode",
        choices=("ar", "mtp", "parity"),
        default="parity",
    )
    parser.add_argument("--max-new-tokens", type=int, default=16)
    parser.add_argument("--depth", type=int)
    cli = parser.parse_args()

    runtime = WispMTPReferenceRuntime.load(cli.package)
    if cli.mode == "ar":
        result: dict[str, Any] = runtime.decode_ar(
            cli.prompt_ids,
            cli.max_new_tokens,
        ).to_dict()
    elif cli.mode == "mtp":
        result = runtime.decode_mtp(
            cli.prompt_ids,
            cli.max_new_tokens,
            cli.depth,
        ).to_dict()
    else:
        result = runtime.verify_greedy_parity(
            cli.prompt_ids,
            cli.max_new_tokens,
            cli.depth,
        )
    output = {
        "package": runtime.package_summary(),
        "result": result,
    }
    print(json.dumps(output, indent=2, sort_keys=True, allow_nan=False))


if __name__ == "__main__":
    main()