spb/wpai0001_uqo

1"""MLX backend (Tier ≥ 1): residual-stream capture and intervention by wrapping decoder layers.23mlx_lm models expose ``model.model.layers`` (a Python list of transformer blocks called as4``layer(h, mask, cache)``). We replace each block by a thin wrapper that calls the original block and then5applies the project's hook contract at ``resid_post.{l}``. No-op guarantee: when the manager is disabled6or has no interventions for that layer, the wrapper returns the block's output array unchanged (same7object), which is verified bitwise by ``tests/integration/test_mlx_noop.py`` when MLX is available.89Only what MLX actually exposes is used: block inputs/outputs (residual stream). Per-head attention10outputs are NOT exposed without re-implementing attention; head-level interventions are therefore11reported as unsupported for the MLX backend.12"""13from __future__ import annotations1415from dataclasses import dataclass, field16from typing import Any, Callable1718try:19    import mlx.core as mx20    import mlx.nn as nn2122    MLX_AVAILABLE = True23except Exception:  # noqa: BLE00124    MLX_AVAILABLE = False252627@dataclass28class MLXHookContext:29    name: str30    extra: dict = field(default_factory=dict)313233class MLXHookManager:34    def __init__(self) -> None:35        self.enabled = True36        self._hooks: dict[str, list[Callable]] = {}37        self.calls = 03839    def add(self, name: str, fn: Callable) -> None:40        self._hooks.setdefault(name, []).append(fn)4142    def remove(self, name: str, fn: Callable | None = None) -> None:43        if fn is None:44            self._hooks.pop(name, None)45        elif name in self._hooks:46            self._hooks[name] = [f for f in self._hooks[name] if f is not fn]4748    def clear(self) -> None:49        self._hooks.clear()5051    def __call__(self, name: str, x: Any, ctx: MLXHookContext | None = None) -> Any:52        if not self.enabled:53            return x54        fns = self._hooks.get(name)55        if not fns:56            return x57        self.calls += 158        c = ctx or MLXHookContext(name)59        c.name = name60        for fn in fns:61            x = fn(x, c)62        return x636465if MLX_AVAILABLE:6667    class WrappedBlock(nn.Module):68        def __init__(self, block: nn.Module, layer: int, hooks: MLXHookManager) -> None:69            super().__init__()70            self.block = block71            self.layer = layer72            self._hooks = hooks  # not a module attribute of the parameter tree7374        def __getattr__(self, key):  # delegate block attributes (e.g. use_sliding) read by the model wrapper75            try:76                return super().__getattr__(key)77            except AttributeError:78                if key in ("block", "layer", "_hooks"):79                    raise80                return getattr(super().__getattr__("block"), key)8182        def __call__(self, x, mask=None, cache=None):83            x = self._hooks(f"resid_pre.{self.layer}", x)84            out = self.block(x, mask, cache)85            return self._hooks(f"resid_post.{self.layer}", out)8687    class MLXCapture:88        """Record the last-token (or all-token) activation of selected layers as numpy arrays."""8990        def __init__(self, positions: str | list[int] = "last") -> None:91            self.positions = positions92            self.store: dict[str, Any] = {}9394        def __call__(self, x, ctx):95            if self.positions == "last":96                self.store[ctx.name] = x[:, -1, :]97            elif self.positions == "all":98                self.store[ctx.name] = x99            else:100                self.store[ctx.name] = x[:, self.positions, :]101            return x102103    class MLXAddAtLast:104        """x[:, -1, :] += alpha * v  (inference-time steering at the last position; single-token)."""105106        def __init__(self, v, alpha: float = 1.0, positions: str = "last") -> None:107            self.v, self.alpha, self.positions = v, alpha, positions108109        def __call__(self, x, ctx):110            v = mx.array(self.v).astype(x.dtype) * self.alpha111            if self.positions == "last":112                delta = mx.zeros_like(x)113                delta[:, -1, :] = v114                return x + delta115            return x + v  # broadcast over all positions116117    class MLXReplaceAtLast:118        def __init__(self, rows, alpha: float = 1.0) -> None:119            self.rows, self.alpha = rows, alpha120121        def __call__(self, x, ctx):122            r = mx.array(self.rows).astype(x.dtype)123            cur = x[:, -1, :]124            new = cur + self.alpha * (r - cur)125            out = mx.array(x)126            out[:, -1, :] = new127            return out128129    def instrument(model) -> MLXHookManager:130        """Wrap all decoder layers in place; idempotent. Returns the hook manager."""131        inner = model.model132        if getattr(inner, "_wp_hooks", None) is not None:133            return inner._wp_hooks134        hooks = MLXHookManager()135        inner.layers = [WrappedBlock(b, i, hooks) for i, b in enumerate(inner.layers)]136        inner._wp_hooks = hooks137        return hooks138139    def last_token_logits(model, tokens: list[list[int]]):140        """Right-aligned batch? mlx_lm uses causal masks without padding support in this helper: one prompt per call."""141        out = []142        for t in tokens:143            logits = model(mx.array([t]))144            out.append(logits[0, -1, :])145        return out146