Laya-MLX: нативный рантайм typed decisions на Apple Silicon

Laya-MLX — это независимый MLX-порт моделей семейства Laya, которые отвечают не текстом, а структурированными решениями: выбором варианта, оценкой по рубрике или вероятностью утверждения. Репозиторий позиционируется как «open-weight typed decisions, работающие нативно на Apple Silicon»: медианное время короткого английского решения — 13,4 мс, для мультиязычного чекпоинта — 7,4 мс, при нуле выходных токенов. Никакого PyTorch, Transformers runtime и облачных API — только локальный inference на MLX.

Что такое typed decisions

Во многих сценариях программе нужен не текст, а конкретный ответ: куда направить обращение, какую оценку поставить по рубрике, верно ли утверждение. Laya решает такие ограниченные вопросы одним двунаправленным проходом, без поколения токенов и без сгенерированного JSON:

state + typed question → bidirectional encoder → decision heads → probabilities

Поддерживается три типа вопросов:

• choice — вероятности по именованным вариантам;
• score — вероятности по уровням рубрики и ожидаемая оценка;
• noul — P(true) для произвольного утверждения.

Строки вопросов батчируются независимо. Скрытые состояния двунаправленного энкодера зависят и от состояния, и от вопроса: рантайм не заявляет, что состояние кодируется один раз и переиспользуется для любых вопросов.

Производительность на M3 Max

Замеры на M3 Max (40 GPU-ядер, 128 ГБ памяти, FP16, end-to-end):

В тайминг входят подготовка промпта, токенизация, тензоры, синхронный inference, калибровка и форматирование результата; загрузка модели не учитывается. Полная методика и все выборки — в BENCHMARKS.md.

Fidelity порта: все три чекпоинта совпали с upstream по выбранному ответу на 63/63 валидационных вопросах в FP32 и FP16 — 378 из 378 сравнений. Каждая конфигурация прошла 100 повторных детерминированных вызовов без роста активной памяти.

Поддерживаемые чекпоинты

Готовые FP16-чекпоинты опубликованы на Hugging Face: aac6fef/laya-mlx, aac6fef/laya-multilingual-mlx, aac6fef/laya-typed-decisions-mlx. Все 36 опубликованных файлов прошли строгую проверку контрольных сумм. Репозиторий даёт inference и конвертацию; обучение RLCD и fine-tuning остаются в upstream-проекте. Это независимый порт, а не официальный релиз Convai Innovations.

Установка и Python API

pip install laya-mlx

import laya_mlx as laya

agent = laya.load("aac6fef/laya-mlx")
result = agent.predict(
    "I was billed twice. Please refund the duplicate.",
    {
        "department": {
            "type": "choice",
            "instructions": "Who should handle this?",
            "criteria": ["billing", "technical", "sales"],
        }
    },
)
print(result["answers"]["department"])

Требуется Apple Silicon, Python 3.11+, macOS 14+. При первой загрузке чекпоинт скачивается, дальше inference полностью локальный. system_one — алиас для predict. Состояния бывают текстом, JSON-словарями или списком сообщений. По умолчанию FP16; dtype="float32" даёт более близкую численную согласованность. batch_size=16 ограничивает число вопросов на forward-проход, device="gpu" или device="cpu" выбирает устройство явно.

Для повторяющихся нагрузок можно включить compile=True, pad_to_multiple=16 и cache_prompts=True при загрузке Agent. Префиксный кэш ограничен 128 вопросами; все три опции по умолчанию выключены, у компиляции есть стоимость первого вызова.

При загрузке проверяются имена и формы всех параметров; паттерны внимания ModernBERT, границы sliding window, базы RoPE и нормализация первого слоя сохраняются. Сырые значения калибровочных температур доступны как agent.temperature_raw: при загрузке температуры зажимаются в [0.5, 5.0], чтобы испорченный бакет не превращал подбрасывание монеты в «почти достоверный» ответ.

Роутинг по языку

Из коробки есть Router и пресеты вопросов, адаптированные из upstream:

from laya_mlx import Router, triage_questions

router = Router(dtype="float16", max_loaded=2)
result = router.predict({"message": "Facture prélevée deux fois, remboursez s'il vous plaît."}, triage_questions())
print(result["routing"])  # multilingual

Router(preload=True) держит все три чекпоинта в памяти; жизненный цикл моделей защищён re-entrant lock, поэтому потоки разделяют одну загруженную модель, а сами вызовы не сериализуются. Неопознанные языки с латиницей (румынский, польский, чешский, турецкий...) уходят в мультиязычный чекпоинт, а не считаются английским. detect_language(state) отдаёт доказательства: language_undecided и diacritic_rate.

Shortlisting для больших наборов вариантов

Варианты choice делят один токен-бюджет, поэтому сотни меток оставляют по несколько токенов на каждую. predict_shortlist эмбеддит состояние и каждую метку, оставляет топ-K по косинусной близости и запускает один predict по сокращённому набору. Это opt-in: обычный Agent.predict по-прежнему оценивает все переданные критерии. Отдельный bi-encoder обычно shortlist'ит лучше, чем собственный энкодер decision-чекпоинта.

CLI и экспорт

uv run laya-mlx predict \
  --model aac6fef/laya-mlx \
  --state-file examples/state.json \
  --questions examples/questions.json

Конвертация чекпоинта в MLX:

uv run laya-mlx convert \
  --model convaiinnovations/laya \
  --dtype float16 \
  --output models/laya-mlx-fp16

Экспорт содержит model.safetensors, конфигурации энкодера и агента, токенизатор и mlx_config.json. Это конвертация имён параметров и dtype, а не квантизация или переобучение; существующие каталоги не перезаписываются.

Демо, тесты и исследования

Терминальная игра «Змейка» (pip install 'laya-mlx[demo]', затем laya-snake) принимает решение Laya на каждый ход. В режиме --optimize --max-speed на M3 Max зафиксировано 75,40 ходов/с на 2400 ходов, ноль смертей и 2 видимых вмешательства safety-слоя — примерно на 6,5% быстрее eager-контроля в том же прогоне.

Тесты покрывают сравнения с Transformers и upstream decision head, валидацию на реальных чекпоинтах и бенчмарки в свежем процессе на каждый бэкенд. В репозитории лежат отдельные отчёты по исследованию производительности: начальное исследование, математический разбор варианта 10× и инженерный разбор. Текущие результаты не подтверждают универсального ещё одного 10× ускорения на тех же чекпоинтах: в отдельных случаях даётся около 1,03–1,08× по медиане.

Лицензия

Apache-2.0; см. LICENSE и NOTICE. Laya и предобученные веса принадлежат Convai Innovations и авторам upstream. Построение промпта, форматирование вывода, языковой роутинг и пресеты адаптированы из NandhaKishorM/laya; архитектура нейросети реализована заново на MLX по образцу Laya и Hugging Face ModernBERT.