Skip to content

Add MLX backend for raw / L0 / L1 - #15

Open
TerenceLiu98 wants to merge 1 commit into
nokia-applied-research:mainfrom
TerenceLiu98:main
Open

TerenceLiu98 wants to merge 1 commit into
nokia-applied-research:mainfrom
TerenceLiu98:main

Conversation

@TerenceLiu98

Copy link
Copy Markdown
## Summary

This PR adds an optional MLX-LM backend for running AnyJev on Apple Silicon.

The backend supports raw, L0 and L1 decisions through direct next-token logits. L2 is intentionally not supported because the backend does not expose a stable hidden-state interface.

## Changes

- Add `MLXBackend` under `anyjev/backends/mlx.py`
- Add the optional `mlx` dependency extra
- Add `scripts/mlx_parity.py` for comparison with `HFBackend`
- Add unit tests covering:
  - batched inference
  - right-padding
  - last valid token position
  - restricted log-probability extraction
  - invalid prompt inputs
- Add an `@pytest.mark.engine` smoke test for a real MLX-LM model
- Document installation and usage in `README.md`
- Record MLX and MLX-LM licensing in `THIRD_PARTY.md`
- Add a changelog entry

## Usage

```bash
pip install -e ".[mlx]"
from anyjev import Decider
from anyjev.backends.mlx import MLXBackend

backend = MLXBackend(
    "Qwen/Qwen3-0.6B-MLX-4bit"
)
decider = Decider(backend, level="L0")

Validation

  • ruff check anyjev bench demo scripts space tests
  • pytest -q -m 'not engine' --ignore=tests/test_truncate_and_vllm.py
  • pytest -q -m engine tests/test_mlx_backend.py

The MLX engine test is skipped when MLX, MLX-LM, Metal, or the configured model is unavailable.

No benchmark numbers are added in this PR. The parity script reports differences when run with matching Transformers and MLX models.

@MorrisZJ

MorrisZJ commented Oct 7, 2026

Copy link
Copy Markdown
Collaborator

Thanks @TerenceLiu98, and sorry for the wait; 0.3.0 landed in between. We ran this on Apple Silicon (M5, mlx 0.32.3, mlx-lm 0.32.0, Qwen2.5-0.5B-Instruct): your unit tests and the engine smoke test pass; scripts/mlx_parity.py against transformers on CPU (fp32) gives 15/15 argmax agreement (max |Δ| 0.17 restricted log-softmax, MLX in bf16); and L0 decisions through Decider match on 23/24 choice, yes/no and score items (max |Δp| 0.03; the one flip is a near-tie). So the readout is right. A few changes before we merge:

  1. Logits at the last position only. self.model(batch) materialises full-vocabulary logits at every position before the last one is selected. Measured at batch 8 × ~1,800 tokens: peak 6.29 GB, against 1.95 GB when the backbone runs first and the LM head is applied to the last position only (h = model.model(x)[rows, lengths - 1], then model.lm_head(h), or model.model.embed_tokens.as_linear(h) for tied embeddings). The output is bit-identical and about 30% faster, and the gap grows with the model and the prompt, which matters on 8–16 GB Macs. Please fall back to the full forward if a model exposes neither.
  2. The parity script can't run on a Mac. HFBackend defaults to device="cuda"; please add --hf-device (default cpu, with float32).
  3. README: 0.3.0 rewrote it and the section this edits no longer exists. Please drop the README change; the docstring is enough.
  4. 0.3.0 removed L1: "raw / L0" in the title, docstring and CHANGELOG.
  5. Rebase onto main (CHANGELOG.md, pyproject.toml, THIRD_PARTY.md), and note the mlx / mlx-lm versions you checked in the docstring.

We'll rerun the same checks and merge.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants