From 40c7b1c951f1c5fe48fda95d3bea794ef727af63 Mon Sep 17 00:00:00 2001 From: AndreSlavescu Date: Fri, 31 Jul 2026 03:37:37 -0700 Subject: [PATCH] Add metile.compile(model): one call, structural detection, verified before it is kept The integration worked but nobody could reach it. Accelerating a model meant importing from metile.integrations.mlx_lm, choosing among eleven keyword flags, and knowing which combination was safe for the architecture in hand. This is the Liger-shaped entry point: import metile print(metile.compile(model)) Two changes underneath it, both from failures this project has had rather than from taste. Architectures are now matched by structure as well as by name. The patcher gated on a list of module and class names, which only ever covers what someone remembered to add: Qwen3.5, Qwen3.6 and Qwen3-VL were all excluded by it, and their equivalence tests reported skips that read like passes. A class carrying gate_proj, up_proj and down_proj is now a candidate whether or not it has been seen, which is what makes an unlisted model work today. Structure alone is not enough to act on, so compile verifies. A class can have the parts of a gated MLP and still scale the product or use a different activation, presenting identically. So it runs the model before and after, compares decode-step logits, and keeps only what reproduces MLX. Decode steps rather than prefill, because attention only engages at query length one and a prefill comparison reports agreement while never running the kernel. When the full set disagrees it bisects per feature instead of reverting everything, which is what keeps three of four features on Llama-3.2-1B rather than none. Comparison is exact by default and the report says what that costs. Llama-3.2-1B's quantized_mlp moves a logit by 0.035 against a magnitude near 20, 2.3e-3 relative -- a summation-order difference where meTile measured as the more accurate side, 4.10 against MLX's 18.05 versus a float32 reference. Declining it by default is the conservative call, and the report names the feature, the absolute and relative size, and that 5e-3 tolerance would keep it, so the trade is visible before it is taken rather than discovered later. The report is falsy when nothing was replaced. That is the point of it. The dangerous outcome here is not a crash but a silent no-op, and sixteen tests in this repository reported "skipped" for models meTile was not touching for as long as it took someone to read the skip list. Verified on three real checkpoints: Qwen2.5-1.5B and Qwen3.5-4B take all four features with exactly matching logits, and Llama-3.2-1B takes three with the fourth declined and explained. 693 pass. Lint and vulture clean. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 45 ++++++ metile/__init__.py | 3 + metile/compile.py | 254 ++++++++++++++++++++++++++++++++++ metile/integrations/mlx_lm.py | 65 ++++++++- tests/test_compile.py | 141 +++++++++++++++++++ 5 files changed, 505 insertions(+), 3 deletions(-) create mode 100644 metile/compile.py create mode 100644 tests/test_compile.py diff --git a/README.md b/README.md index 0795e45..833ca0f 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,51 @@ You write the obvious three passes. The compiler notices the first two can be me one and rewrites them. That reads the input twice instead of three times and runs 1.28x faster, but only after checking that the merge is algebraically valid. +## Accelerating an MLX-LM model + +One call. It patches what it recognises, checks the result against the unpatched model, and +tells you what it did. + +```python +import metile +from mlx_lm import load + +model, tokenizer = load("mlx-community/Qwen2.5-1.5B-Instruct-4bit") +print(metile.compile(model)) +``` + +``` +meTile on qwen2 + accelerating: attention, rms_norm, graph_fusion, quantized_mlp + surfaces replaced: mlp, input_layernorm, post_attention_layernorm, block + verification: logits match MLX exactly +``` + +Architectures are matched by structure rather than by a list of names, so a model with the +usual gated MLP is a candidate whether or not anyone has seen it before. Structure is only a +candidate test though: a class can have `gate_proj`, `up_proj` and `down_proj` and still +combine them differently. So `compile` runs the model before and after, compares the logits, +and keeps only what reproduces MLX. If the whole set disagrees it bisects and keeps the parts +that pass: + +``` +meTile on llama + accelerating: attention, rms_norm, graph_fusion + verification: logits match MLX exactly + declined quantized_mlp: changed the logits by 0.0352, 2.3e-03 relative + -- reduction-order scale, raise tolerance to keep it +``` + +That one is a summation-order difference where meTile is the *more* accurate side, so +`metile.compile(model, tolerance=5e-3)` keeps it. The default is exact because the failure +worth catching -- a structural match whose arithmetic differs -- lands far outside rounding. + +The report is falsy when nothing was replaced, so `if not metile.compile(model)` is a real +check. That matters more than it sounds: the dangerous outcome is not a crash, it is a silent +no-op, and this project shipped one for three model families before anyone read the skip list. + +Call `.restore()` on the report to put MLX-LM's own implementations back. + ## Speed Everything below runs on one **Apple M5 (32 GB, MLX 0.32.0)** and compares meTile against diff --git a/metile/__init__.py b/metile/__init__.py index 5d7647b..ca821d5 100644 --- a/metile/__init__.py +++ b/metile/__init__.py @@ -1,3 +1,4 @@ +from metile.compile import CompileReport, compile from metile.compiler.graph_fusion import FusionTarget, ParallelEpilogueRule, plan_graph_fusion from metile.frontend.autotune import Config, autotune from metile.frontend.kernel import kernel @@ -68,6 +69,7 @@ __all__ = [ "BlockScaledWeight", "Buffer", + "CompileReport", "ComputeGraph", "Config", "FusionTarget", @@ -88,6 +90,7 @@ "cast", "cdiv", "col_major", + "compile", "constexpr", "dot", "exp", diff --git a/metile/compile.py b/metile/compile.py new file mode 100644 index 0000000..b9e6a72 --- /dev/null +++ b/metile/compile.py @@ -0,0 +1,254 @@ +"""One call that accelerates an MLX-LM model and tells you what it did. + + import metile + model, tokenizer = mlx_lm.load("mlx-community/Qwen2.5-1.5B-Instruct-4bit") + print(metile.compile(model)) + +Two things distinguish this from calling the patcher directly, and both come from failures this project +has actually had. + +**It recognises architectures nobody enumerated.** The underlying patcher gates on a list of module and +class names, which only ever covers what someone remembered to add. Qwen3.5, Qwen3.6 and Qwen3-VL were all +excluded by that list, silently, for as long as it took to notice. `compile` also admits classes by +structure, so a model with the usual gated MLP is a candidate whether or not it has been seen before. + +**It verifies before it keeps anything.** Structure is a weaker claim than a name: a class with +`gate_proj`, `up_proj` and `down_proj` has the parts of a gated MLP but might scale the product or use a +different activation, and it would present identically. So `compile` runs the model before and after +patching and compares the logits, feature by feature if the whole set disagrees, and keeps only what +reproduces MLX's output. A model it cannot verify runs unpatched rather than wrong. + +The report is the third piece. The dangerous outcome here is not a crash, it is a silent no-op: sixteen +equivalence tests in this project spent weeks reporting "skipped" for models where nothing was being +patched, which reads like success in a summary. `compile` returns a report that says what it replaced and +what it declined, and it is falsy when it changed nothing. + +Verification costs a few forward passes. Pass `verify=False` to skip it when the architecture is already +covered by the model matrix, and expect no protection against a structural false positive if you do. +""" + +from dataclasses import dataclass, field + +# Feature flags the patcher exposes that `compile` turns on, in the order they are tried when a combined +# verification fails and the set has to be bisected. +FEATURES = ("attention", "rms_norm", "graph_fusion", "quantized_mlp") + +# Logits are compared exactly by default, because the failure this exists to catch -- a class that has the +# parts of a gated MLP but combines them differently -- lands far outside any rounding. +# +# Exactness does cost something real, and the report says so rather than hiding it. meTile's kernels are +# bit-exact with MLX wherever the arithmetic matches, but a few are not: they sum in a different order, and +# in the two cases checked against a float32 reference meTile was the *more* accurate side, by 4.4x on f16 +# SwiGLU and 1.6x on f16 RMSNorm. Llama-3.2-1B is one of them, where quantized_mlp moves a logit by 0.035 +# against a magnitude near 20 -- about 2e-3 relative, and declined under the default. +# +# So `tolerance` is relative to the largest reference logit and defaults to exact. Around 1e-3 is the scale +# of a reduction-order difference; anything much larger is a different computation, not a different order. +TOLERANCE = 0.0 +REORDERING_SCALE = 5e-3 + + +@dataclass +class CompileReport: + """What `compile` replaced, what it declined, and why.""" + + model: str = "unknown" + features: tuple = () + surfaces: tuple = () + verified: bool | None = None + difference: float | None = None + relative: float | None = None + declined: dict = field(default_factory=dict) + handle: object = None + + def __bool__(self): + """False when nothing was replaced, so `if not metile.compile(model)` is a usable check.""" + return bool(self.features) + + def restore(self): + """Put back MLX-LM's own implementations.""" + if self.handle is not None: + self.handle.__exit__(None, None, None) + self.handle = None + + def __str__(self): + lines = [f"meTile on {self.model}"] + if self.features: + lines.append(f" accelerating: {', '.join(self.features)}") + lines.append(f" surfaces replaced: {', '.join(self.surfaces) or 'none reported'}") + else: + lines.append(" accelerating: nothing -- this model runs entirely on MLX") + if self.verified is None: + lines.append(" verification: skipped (verify=False)") + elif self.verified and not self.difference: + lines.append(" verification: logits match MLX exactly") + elif self.verified: + lines.append( + f" verification: within tolerance, logits differ by {self.difference:g}" + + (f" ({self.relative:.1e} relative)" if self.relative else "") + ) + else: + lines.append(" verification: FAILED, everything reverted") + for feature, reason in sorted(self.declined.items()): + lines.append(f" declined {feature}: {reason}") + return "\n".join(lines) + + +def _decode_logits(model, tokens, steps=4): + """Prefill, take a few decode steps, return the last step's logits. + + Decode steps rather than prefill, because meTile's attention only engages at query length one. A + prefill-only comparison reports agreement while never running the kernel, which is true and + meaningless; this project measured everything as exact until that was noticed. + """ + import mlx.core as mx + from mlx_lm.models.cache import make_prompt_cache + + cache = make_prompt_cache(model) + out = model(tokens, cache=cache) + mx.eval(out) + following = mx.argmax(out[:, -1, :], axis=-1) + for _ in range(steps): + out = model(following[None, :], cache=cache) + mx.eval(out) + following = mx.argmax(out[:, -1, :], axis=-1) + return out[:, -1, :].astype(mx.float32) + + +def _probe_tokens(model): + """A short token sequence valid for this model's vocabulary.""" + import mlx.core as mx + + vocabulary = getattr(model, "vocab_size", None) or 32000 + return mx.array([[index % max(vocabulary - 1, 1) + 1 for index in range(8)]]) + + +def _surfaces(model): + """Names of the layer attributes meTile currently has replacements bound to.""" + from metile.integrations.mlx_lm import _model_layers + + watched = ("mlp", "self_attn", "linear_attn", "input_layernorm", "post_attention_layernorm") + found = [] + for layer in _model_layers(model): + for name in watched: + member = getattr(layer, name, None) + if member is None or name in found: + continue + implementation = type(member).__call__ + if getattr(implementation, "_metile_original", None) is not None: + found.append(name) + block = type(layer).__call__ + if getattr(block, "_metile_original", None) is not None and "block" not in found: + found.append("block") + return tuple(found) + + +def compile(model, *, verify=True, features=FEATURES, tolerance=TOLERANCE): + """Accelerate an MLX-LM model in place and report what changed. + + Returns a `CompileReport`, which is falsy when nothing was replaced. Call `.restore()` on it to put + MLX-LM's implementations back. + + `tolerance` is relative to the largest reference logit and defaults to exact. Raising it to about + 5e-3 admits differences at the scale of a summation reorder, which is where meTile's remaining + divergences from MLX sit and where it measured as the more accurate side; the report names any + feature it declined and by how much, so the trade is visible before it is taken. + """ + from metile.integrations.mlx_lm import apply_metile_to_mlx_lm + + if model is None or not callable(model): + raise TypeError("compile expects a loaded MLX-LM model") + + name = type(getattr(model, "model", model)).__module__.rsplit(".", 1)[-1] + requested = tuple(feature for feature in FEATURES if feature in features) + declined = {} + + reference = None + tokens = None + magnitude = 1.0 + if verify: + try: + tokens = _probe_tokens(model) + reference = _decode_logits(model, tokens) + magnitude = max( + float( + __import__("mlx.core", fromlist=["core"]) + .max(__import__("mlx.core", fromlist=["core"]).abs(reference)) + .item() + ), + 1e-9, + ) + except Exception as error: + declined["verification"] = f"could not run the model unpatched ({type(error).__name__})" + verify = False + + def attempt(selected): + """Patch with `selected` enabled and return the handle, or None if it verifies wrong.""" + flags = {feature: feature in selected for feature in FEATURES} + handle = apply_metile_to_mlx_lm(model=model, **flags) + handle.__enter__() + if not verify: + return handle, None + import mlx.core as mx + + try: + difference = float(mx.max(mx.abs(_decode_logits(model, tokens) - reference)).item()) + except Exception: + handle.__exit__(None, None, None) + return None, None + if difference > tolerance * magnitude: + handle.__exit__(None, None, None) + return None, difference + return handle, difference + + handle, difference = attempt(requested) + kept = requested + if verify and handle is None: + # The combined set disagrees, so find which parts do. Bisecting per feature beats reverting + # everything: usually one surface is at fault and the rest reproduce MLX exactly. + declined["combined"] = "the full feature set changed the logits; bisected" + kept = [] + for feature in requested: + trial, trial_difference = attempt([feature]) + if trial is None: + if trial_difference: + relative = trial_difference / magnitude + scale = ( + "reduction-order scale, raise tolerance to keep it" + if relative <= REORDERING_SCALE + else "far beyond rounding, so a different computation" + ) + declined[feature] = ( + f"changed the logits by {trial_difference:g}, " + f"{relative:.1e} relative -- {scale}" + ) + else: + declined[feature] = "could not run" + continue + trial.__exit__(None, None, None) + kept.append(feature) + if kept: + handle, difference = attempt(kept) + if handle is None: + kept = [] + + if handle is None: + return CompileReport(model=name, verified=False if verify else None, declined=declined) + + surfaces = _surfaces(model) + if not surfaces: + # Patching a class nothing in this model uses is the silent no-op this report exists to expose. + handle.__exit__(None, None, None) + declined["surfaces"] = "no layer in this model uses a class meTile can replace" + return CompileReport(model=name, verified=None, declined=declined) + + return CompileReport( + model=name, + features=tuple(kept), + surfaces=surfaces, + verified=None if not verify else True, + difference=difference, + relative=(difference / magnitude) if (verify and difference is not None) else None, + declined=declined, + handle=handle, + ) diff --git a/metile/integrations/mlx_lm.py b/metile/integrations/mlx_lm.py index 1af0251..f57a1b0 100644 --- a/metile/integrations/mlx_lm.py +++ b/metile/integrations/mlx_lm.py @@ -173,9 +173,68 @@ def _attention_module(block): return None -def _recognised(cls, registry): - """Whether meTile is allowed to replace this class's __call__.""" - return (cls.__module__, cls.__name__) in registry +# Attributes a class must carry to be a candidate for each replacement. Structure rather than name, +# because a name list only ever covers the architectures someone remembered: Qwen3.5, Qwen3.6 and +# Qwen3-VL were all excluded by one until it was noticed, and their equivalence tests reported skips that +# read like passes. +# +# Structure is a weaker claim than the name list it supplements. `gate_proj`/`up_proj`/`down_proj` says a +# class has the parts of a gated MLP, not that its __call__ combines them the way meTile's replacement +# does -- a model scaling the product, or applying a different activation, presents identically. So a +# structural match admits a candidate and nothing more; `metile.compile` runs the model and compares +# against the unpatched result before keeping any of it. +_GATED_MLP_ATTRIBUTES = ("gate_proj", "up_proj", "down_proj") +_FUSED_BLOCK_ATTRIBUTES = ("input_layernorm", "post_attention_layernorm", "mlp") + +_STRUCTURE = { + id(_GATED_MLP_CLASSES): _GATED_MLP_ATTRIBUTES, + id(_FUSED_BLOCK_CLASSES): _FUSED_BLOCK_ATTRIBUTES, +} + + +def _structurally_matches(cls, registry): + """Whether a class carries the parts the registry's replacement needs. + + Read off the class, so it works for architectures nobody enumerated. Requires __call__ to be defined + on the class itself rather than inherited, because a class that does not define one has nothing for + meTile to replace and `getattr` would hand back a fresh method-wrapper from the metaclass. + """ + required = _STRUCTURE.get(id(registry)) + if required is None: + return False + if not any("__call__" in vars(klass) for klass in cls.__mro__): + return False + if not all( + hasattr(cls, name) or name in getattr(cls, "__annotations__", {}) for name in required + ): + # Attributes are usually set in __init__ rather than declared, so fall back to the source of + # __call__: a replacement only works if the body actually reaches those names. + source = _call_source(cls) + return source is not None and all(name in source for name in required) + return True + + +def _call_source(cls): + import inspect as _inspect + + for klass in cls.__mro__: + if "__call__" in vars(klass): + try: + return _inspect.getsource(vars(klass)["__call__"]) + except (OSError, TypeError): + return None + return None + + +def _recognised(cls, registry, structural=True): + """Whether meTile is allowed to replace this class's __call__. + + The named pairs are the combinations whose arithmetic has been checked against MLX in the model + matrix. Structural matches are candidates that `metile.compile` verifies at patch time. + """ + if (cls.__module__, cls.__name__) in registry: + return True + return structural and _structurally_matches(cls, registry) def _registry_classes(registry): diff --git a/tests/test_compile.py b/tests/test_compile.py new file mode 100644 index 0000000..e4847c9 --- /dev/null +++ b/tests/test_compile.py @@ -0,0 +1,141 @@ +"""The one-call entry point, and the failure it is built to make impossible. + +`metile.compile(model)` is the API people actually use, so the thing it must never do is quietly change +nothing and look like it worked. That failure has a track record here: sixteen equivalence tests reported +"skipped" for models meTile was not touching, which reads like success in a summary, for as long as it took +someone to read the skip list. + +So the report is falsy when nothing was replaced, it names what it declined and by how much, and it +verifies against the unpatched model before keeping anything. +""" + +import pytest + +import metile +from metile.compile import CompileReport + + +def test_a_report_that_replaced_nothing_is_falsy(): + """`if not metile.compile(model)` has to be a usable check, or silence looks like success.""" + assert not CompileReport(model="whatever") + assert CompileReport(model="whatever", features=("attention",)) + + +def test_the_report_says_plainly_when_it_did_nothing(): + """A summary someone skims must not read as success.""" + text = str(CompileReport(model="exotic")) + assert "nothing" in text + assert "runs entirely on MLX" in text + + +def test_the_report_distinguishes_exact_from_within_tolerance(): + """Reporting a tolerated difference as "exact" would hide the trade the caller opted into.""" + exact = str(CompileReport(model="m", features=("attention",), verified=True, difference=0.0)) + assert "match MLX exactly" in exact + + tolerated = str( + CompileReport( + model="m", + features=("attention",), + verified=True, + difference=0.035, + relative=2.3e-3, + ) + ) + assert "within tolerance" in tolerated + assert "exactly" not in tolerated + + +def test_compile_rejects_something_that_is_not_a_model(): + for value in (None, 42, "a string"): + with pytest.raises(TypeError, match="MLX-LM model"): + metile.compile(value) + + +def test_structural_detection_admits_unlisted_architectures(): + """The generalisation, checked on a class that is deliberately not in the name list. + + A name list only covers what someone remembered, and this one excluded Qwen3.5, Qwen3.6 and Qwen3-VL + without saying so. Structure is what makes an unseen model a candidate; verification is what makes that + safe, and the two are separate steps on purpose. + """ + pytest.importorskip("mlx_lm") + from mlx_lm.models import gemma2, llama + + from metile.integrations.mlx_lm import ( + _FUSED_BLOCK_CLASSES, + _GATED_MLP_CLASSES, + _recognised, + ) + + assert _recognised(llama.MLP, _GATED_MLP_CLASSES) + assert _recognised(llama.MLP, _GATED_MLP_CLASSES, structural=False) + + # Not in the list, and admitted anyway because it has the parts. + assert _recognised(gemma2.MLP, _GATED_MLP_CLASSES) + assert not _recognised(gemma2.MLP, _GATED_MLP_CLASSES, structural=False) + + # Structure has to exclude as well as admit, or it is not doing any work. + assert not _recognised(llama.Attention, _GATED_MLP_CLASSES) + assert not _recognised(llama.MLP, _FUSED_BLOCK_CLASSES) + + +def test_structural_detection_needs_a_call_to_replace(): + """A class that defines no __call__ has nothing to swap, and getattr would not say so. + + `getattr(cls, "__call__")` on such a class resolves through the metaclass and returns a fresh + method-wrapper every time, which reads as "present". This project has already been caught by that once. + """ + from metile.integrations.mlx_lm import _GATED_MLP_CLASSES, _structurally_matches + + class HasThePartsButNoCall: + gate_proj = up_proj = down_proj = staticmethod(lambda x: x) + + assert not _structurally_matches(HasThePartsButNoCall, _GATED_MLP_CLASSES) + + +@pytest.mark.slow +def test_compile_on_a_real_model_verifies_and_reports(): + """End to end, on the smallest checkpoint the cache is likely to hold.""" + from pathlib import Path + + pytest.importorskip("mlx_lm") + repo = "mlx-community/Qwen2.5-0.5B-Instruct-4bit" + cache = Path.home() / ".cache/huggingface/hub" / f"models--{repo.replace('/', '--')}" + if not cache.exists(): + pytest.skip(f"{repo} is not in the local cache") + + from mlx_lm import load + + model, _ = load(repo) + report = metile.compile(model) + try: + assert report, f"compile replaced nothing:\n{report}" + assert report.verified is True + assert report.surfaces, "features were kept but no layer reports a replacement" + assert report.difference == 0.0 + finally: + report.restore() + + +@pytest.mark.slow +def test_restore_puts_back_what_it_replaced(): + """A patch that cannot be undone is a patch nobody can measure against.""" + from pathlib import Path + + pytest.importorskip("mlx_lm") + repo = "mlx-community/Qwen2.5-0.5B-Instruct-4bit" + cache = Path.home() / ".cache/huggingface/hub" / f"models--{repo.replace('/', '--')}" + if not cache.exists(): + pytest.skip(f"{repo} is not in the local cache") + + from mlx_lm import load + + from metile.compile import _surfaces + + model, _ = load(repo) + assert not _surfaces(model) + report = metile.compile(model, verify=False) + assert _surfaces(model) + report.restore() + assert not _surfaces(model), "restore left meTile implementations bound"