optimized_runtime.py

"""Text-only loader for the Bonsai 2 MLX pack.

This keeps the language model numerically identical to the pack's bundled loader,
but does not instantiate or evaluate the 0.92 GB vision tower for text-only use.
"""

from __future__ import annotations

import json
import math
import os
import sys
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
from mlx_lm.models.qwen3_5 import TextModel, TextModelArgs
from mlx_lm.models.activations import swiglu
from mlx_lm.models.base import scaled_dot_product_attention


def _touch_placeholder_cache(cache):
    """Give pruned-layer ArraysCache entries serializable, inert state."""
    if cache is not None and cache[0] is None:
        cache[0] = mx.zeros((1,), dtype=mx.uint8)
        cache[1] = mx.zeros((1,), dtype=mx.uint8)


class SkippedDecoderLayer(nn.Module):
    """Identity replacement used only by the explicitly enabled Turbo mode."""

    def __init__(self, is_linear, scale=1.0):
        super().__init__()
        self.original_is_linear = is_linear
        # Identity layers need only a harmless cache placeholder. ArraysCache
        # tolerates an empty state; KVCache.state assumes allocated keys.
        self.is_linear = True
        self.scale = float(scale)

    def __call__(self, x, mask=None, cache=None):
        _touch_placeholder_cache(cache)
        return x if self.scale == 1.0 else x * self.scale


class AffineSkippedDecoderLayer(nn.Module):
    """Cheap calibrated diagonal approximation for an omitted decoder block."""

    def __init__(self, is_linear, scale, bias):
        super().__init__()
        self.original_is_linear = is_linear
        self.is_linear = True
        self.scale = scale
        self.bias = bias

    def __call__(self, x, mask=None, cache=None):
        _touch_placeholder_cache(cache)
        return x * self.scale + self.bias


class SkippedMLP(nn.Module):
    """Zero residual used to omit only a decoder block's MLP branch."""

    def __call__(self, x):
        return mx.zeros_like(x)


class AffineSkippedMLP(nn.Module):
    """Calibrated diagonal approximation of an MLP residual branch."""

    def __init__(self, scale, bias):
        super().__init__()
        self.scale = scale
        self.bias = bias

    def __call__(self, x):
        return x * self.scale + self.bias


class FusedDecodeProjection(nn.Module):
    """Fuse sibling ternary GEMVs only when autoregressive decode benefits."""

    def __init__(self, projections, fwht):
        super().__init__()
        first = projections[0]
        if any(
            projection.block != first.block or projection.signs is not first.signs
            for projection in projections[1:]
        ):
            raise ValueError("Cannot fuse projections with different transforms")
        self.weight = mx.concatenate([p.weight for p in projections], axis=0)
        self.scales = mx.concatenate([p.scales for p in projections], axis=0)
        self.biases = mx.concatenate([p.biases for p in projections], axis=0)
        self.signs = first.signs
        self.block = first.block
        sizes = [projection.weight.shape[0] for projection in projections]
        self.boundaries = [sum(sizes[: index + 1]) for index in range(len(sizes) - 1)]
        self.slices = []
        start = 0
        for size in sizes:
            self.slices.append(slice(start, start + size))
            start += size
        self._fwht = fwht

    def _matmul(self, x, weight, scales, biases):
        return mx.quantized_matmul(
            x,
            weight,
            scales,
            biases,
            transpose=True,
            group_size=128,
            bits=2,
        )

    def __call__(self, x):
        x = self._fwht(x, self.block, self.signs)
        # One wide GEMV is faster during autoregressive decode. For prompt
        # batches, two narrower QMMs are slightly faster on Apple M4.
        if math.prod(x.shape[:-1]) == 1:
            combined = self._matmul(x, self.weight, self.scales, self.biases)
            return tuple(mx.split(combined, self.boundaries, axis=-1))
        else:
            return tuple(
                self._matmul(
                    x,
                    self.weight[section],
                    self.scales[section],
                    self.biases[section],
                )
                for section in self.slices
            )


class FusedDecodeMLP(nn.Module):
    def __init__(self, projection, down):
        super().__init__()
        self.projection = projection
        self.down_proj = down

    def __call__(self, x):
        gate, up = self.projection(x)
        return self.down_proj(swiglu(gate, up))


class FusedDecodeAttention(nn.Module):
    def __init__(self, attention, projection):
        super().__init__()
        self.projection = projection
        self.num_key_value_heads = attention.num_key_value_heads
        self.num_attention_heads = attention.num_attention_heads
        self.head_dim = attention.head_dim
        self.scale = attention.scale
        self.q_norm = attention.q_norm
        self.k_norm = attention.k_norm
        self.rope = attention.rope
        self.o_proj = attention.o_proj

    def __call__(self, x, mask=None, cache=None):
        batch, length, _ = x.shape
        q_proj_output, keys, values = self.projection(x)
        queries, gate = mx.split(
            q_proj_output.reshape(batch, length, self.num_attention_heads, -1),
            2,
            axis=-1,
        )
        gate = gate.reshape(batch, length, -1)
        queries = self.q_norm(queries).transpose(0, 2, 1, 3)
        keys = self.k_norm(
            keys.reshape(batch, length, self.num_key_value_heads, -1)
        ).transpose(0, 2, 1, 3)
        values = values.reshape(
            batch, length, self.num_key_value_heads, -1
        ).transpose(0, 2, 1, 3)

        if cache is not None:
            queries = self.rope(queries, offset=cache.offset)
            keys = self.rope(keys, offset=cache.offset)
            keys, values = cache.update_and_fetch(keys, values)
        else:
            queries = self.rope(queries)
            keys = self.rope(keys)
        output = scaled_dot_product_attention(
            queries, keys, values, cache=cache, scale=self.scale, mask=mask
        )
        output = output.transpose(0, 2, 1, 3).reshape(batch, length, -1)
        return self.o_proj(output * mx.sigmoid(gate))


def load_text_model(directory: str | Path):
    directory = Path(directory)
    config = json.loads((directory / "config.json").read_text())
    if config.get("schema_version") != 2:
        raise ValueError(f"Unsupported schema version: {config.get('schema_version')!r}")
    if config.get("model_type") != "prism_hadamard_qwen35":
        raise ValueError(f"Unsupported model type: {config.get('model_type')!r}")

    runtime = directory / "runtime"
    sys.path.insert(0, str(runtime))
    from runtime import Packed, fwht

    model = TextModel(TextModelArgs.from_dict(config["text_config"]))
    all_weights = mx.load(str(directory / "model.safetensors"))
    prefix = "language_model."
    layout = config.get("optimized_weight_layout", {})
    share_signs = layout.get("sign_sharing") == "by_input_width"
    Packed.share_fwht_activations = bool(layout.get("shared_fwht_activations", False))

    seen = set()
    for record in config["modules"]:
        path = record["path"]
        if path in seen:
            raise ValueError(f"Duplicate packed module: {path}")
        seen.add(path)
        if record["dtype"] != "float16":
            raise ValueError(f"Unsupported activation dtype at {path}")

        parts = path.split(".")
        parent = model
        for part in parts[:-1]:
            parent = parent[int(part)] if part.isdigit() else getattr(parent, part)

        key = prefix + path
        arrays = [all_weights[key + "." + suffix] for suffix in ("weight", "scales", "biases")]
        block = record["block"]
        if block and block not in (512, 1024, 2048, 4096):
            raise ValueError(f"Unsupported Hadamard block size at {path}: {block}")
        signs = all_weights.get(key + ".signs")
        if block and signs is None:
            raise ValueError(f"Missing Hadamard signs at {path}")
        if signs is not None and not mx.all((signs == 1) | (signs == -1)).item():
            raise ValueError(f"Invalid Hadamard signs at {path}")
        setattr(
            parent,
            parts[-1],
            Packed(arrays, block, signs, record["embedding"], mx.float16),
        )

    text_weights = [
        (name[len(prefix) :], value)
        for name, value in all_weights.items()
        if name.startswith(prefix)
    ]
    model.load_weights(text_weights, strict=True)
    del text_weights

    # The checkpoint contains 402 sign tensors but only three unique, byte-identical
    # payloads (one for each input width). Point modules with the same width at one
    # canonical array after load_weights has assigned the checkpoint parameters.
    # This is an exact representation change: ternary weights, scales, biases, and
    # every sign value remain untouched.
    if share_signs:
        canonical_signs = {}
        for record in config["modules"]:
            if not record["block"]:
                continue
            parent = model
            parts = record["path"].split(".")
            for part in parts:
                parent = parent[int(part)] if part.isdigit() else getattr(parent, part)
            signs = parent.signs
            key = (tuple(signs.shape), signs.dtype)
            canonical = canonical_signs.get(key)
            if canonical is None:
                canonical_signs[key] = signs
            else:
                parent.signs = canonical

        expected = layout.get("unique_sign_payloads")
        if expected is not None and len(canonical_signs) != expected:
            raise ValueError(
                f"Expected {expected} canonical sign tensors, found {len(canonical_signs)}"
            )

    fuse_decode = os.environ.get("BONSAI_FUSE_DECODE") == "1" or bool(
        layout.get("fused_decode_projections", False)
    )
    if fuse_decode:
        for layer_index, layer in enumerate(model.model.layers):
            old_mlp = layer.mlp
            gate, up = old_mlp.gate_proj, old_mlp.up_proj
            projection = FusedDecodeProjection((gate, up), fwht)
            fused = FusedDecodeMLP(projection, old_mlp.down_proj)
            # Materialize one layer at a time and release its old gate/up arrays,
            # keeping peak memory bounded while the checkpoint is memory-mapped.
            mx.eval(projection.weight, projection.scales, projection.biases)
            layer.mlp = fused
            stem = f"{prefix}model.layers.{layer_index}.mlp."
            for projection in ("gate_proj", "up_proj"):
                for suffix in ("weight", "scales", "biases", "signs"):
                    all_weights.pop(stem + projection + "." + suffix, None)

            if not layer.is_linear:
                old_attention = layer.self_attn
                qkv = FusedDecodeProjection(
                    (
                        old_attention.q_proj,
                        old_attention.k_proj,
                        old_attention.v_proj,
                    ),
                    fwht,
                )
                mx.eval(qkv.weight, qkv.scales, qkv.biases)
                layer.self_attn = FusedDecodeAttention(old_attention, qkv)
                stem = f"{prefix}model.layers.{layer_index}.self_attn."
                for projection_name in ("q_proj", "k_proj", "v_proj"):
                    for suffix in ("weight", "scales", "biases", "signs"):
                        all_weights.pop(
                            stem + projection_name + "." + suffix, None
                        )

    profile_name = os.environ.get("BONSAI_PROFILE", "").strip()
    profile = {}
    if profile_name:
        profiles = config.get("optimized_runtime_profiles", {})
        if profile_name not in profiles:
            raise ValueError(f"Unknown BONSAI_PROFILE value: {profile_name!r}")
        profile = profiles[profile_name]
        profile_layers = profile.get("skip_layers", [])
        skip_text = ",".join(str(index) for index in profile_layers)
        config["active_runtime_profile"] = profile_name
    else:
        skip_text = os.environ.get("BONSAI_SKIP_LAYERS", "").strip()
    if skip_text:
        skip_layers = sorted({int(value) for value in skip_text.split(",")})
        adapter_path = os.environ.get("BONSAI_SKIP_ADAPTERS", "").strip()
        if not adapter_path and profile.get("skip_adapters"):
            adapter_path = str(directory / profile["skip_adapters"])
        skip_adapters = mx.load(adapter_path) if adapter_path else {}
        scale_text = os.environ.get("BONSAI_SKIP_SCALES", "").strip()
        skip_scales = {}
        if scale_text:
            for item in scale_text.split(","):
                index_text, scale_value = item.split(":", 1)
                skip_scales[int(index_text)] = float(scale_value)
        if any(index < 0 or index >= len(model.model.layers) for index in skip_layers):
            raise ValueError(f"Invalid BONSAI_SKIP_LAYERS value: {skip_text!r}")
        # Layers 0 and 3 own the reference SSM/attention caches used to build
        # masks for the whole stack and therefore cannot be identity-pruned.
        if {0, model.model.fa_idx} & set(skip_layers):
            raise ValueError("Turbo mode cannot skip the reference cache layers")
        for index in skip_layers:
            old_layer = model.model.layers[index]
            scale_key = f"layers.{index}.scale"
            bias_key = f"layers.{index}.bias"
            if adapter_path:
                if scale_key not in skip_adapters or bias_key not in skip_adapters:
                    raise ValueError(f"Missing affine adapter for skipped layer {index}")
                model.model.layers[index] = AffineSkippedDecoderLayer(
                    old_layer.is_linear,
                    skip_adapters[scale_key],
                    skip_adapters[bias_key],
                )
            else:
                model.model.layers[index] = SkippedDecoderLayer(
                    old_layer.is_linear, skip_scales.get(index, 1.0)
                )
        config["active_runtime_skip_layers"] = skip_layers
        if adapter_path:
            config["active_runtime_skip_adapters"] = adapter_path
        if skip_scales:
            config["active_runtime_skip_scales"] = skip_scales

    skip_mlp_text = os.environ.get("BONSAI_SKIP_MLP_LAYERS", "").strip()
    if not skip_mlp_text and profile.get("skip_mlp_layers"):
        skip_mlp_text = ",".join(str(index) for index in profile["skip_mlp_layers"])
    if skip_mlp_text:
        skip_mlp_layers = sorted({int(value) for value in skip_mlp_text.split(",")})
        mlp_adapter_path = os.environ.get("BONSAI_SKIP_MLP_ADAPTERS", "").strip()
        mlp_adapters = mx.load(mlp_adapter_path) if mlp_adapter_path else {}
        if any(index < 0 or index >= len(model.model.layers) for index in skip_mlp_layers):
            raise ValueError(f"Invalid BONSAI_SKIP_MLP_LAYERS value: {skip_mlp_text!r}")
        active_full_skips = set(config.get("active_runtime_skip_layers", []))
        if active_full_skips & set(skip_mlp_layers):
            raise ValueError("A layer cannot be both fully skipped and MLP-skipped")
        for index in skip_mlp_layers:
            scale_key = f"mlp_layers.{index}.scale"
            bias_key = f"mlp_layers.{index}.bias"
            if mlp_adapter_path:
                if scale_key not in mlp_adapters or bias_key not in mlp_adapters:
                    raise ValueError(f"Missing affine MLP adapter for layer {index}")
                model.model.layers[index].mlp = AffineSkippedMLP(
                    mlp_adapters[scale_key], mlp_adapters[bias_key]
                )
            else:
                model.model.layers[index].mlp = SkippedMLP()
        config["active_runtime_skip_mlp_layers"] = skip_mlp_layers
        if mlp_adapter_path:
            config["active_runtime_skip_mlp_adapters"] = mlp_adapter_path

    model.eval()
    mx.eval(model.parameters())
    return model, config