#!/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()