Запуск LLM на Apple Neural Engine: реверс-инжиниринг и запуск Qwen без CPU и GPU

Я продолжаю серию практических исследований по реверс-инжинирингу нейроускорителя Apple Neural Engine (ANE). На этот раз моя цель — разобраться, какими инструкциями и структурами данных нужно снабдить компилятор Apple, чтобы полноценно запустить языковую модель семейства Qwen напрямую на специализированном чипе, полностью исключив нагрузки на CPU и GPU.

Ищем примеры MIL

Разбирать машинный код компилятора вручную — задача чересчур объёмная и нерациональная. Общедоступные примеры промежуточного представления MIL (Model Intermediate Language) в сети довольно однотипны и в большинстве своём цитируют исследования Maderix Claude. Я задался вопросом: где найти действительно валидные и разнообразные образцы кода под чипы Apple Silicon? Ответ очевиден — внутри самой операционной системы macOS.

Выполнив сканирование системного каталога /System/Library на наличие файлов с расширением .mil, я обнаружить порядка 163 файлов. Это полноценные примеры кода MIL, созданные разработчиками Apple. Однако здесь требуется важная оговорка: все они написаны для спецификации CoreML MIL. Не каждый подобный файл удастся скомпилировать и успешно исполнить непосредственно на физическом модуле ANE без предварительной адаптации.

Анализируем находки

Детальное изучение найденных системных файлов позволило выделить несколько ключевых особенностей устройства языка и возможностей компилятора:

• Динамические формы (Dynamic shapes): Поиски аналогов неопределенных размеров (символа ?) привели к аннотации FlexibleShapeInformation. Она содержит блоки DefaultShapes, а также параметры RangeDims и EnumeratedShapes.
• Новые типы данных: В спецификации присутствуют как языковые типы (tuple, list, dict), так и специализированные буферы для работы с тензорами и изображениями (pixel_buffer, tensor_buffer, state<T>).
• Расширенный набор операций: В файлах обнаружилось внушительное количество поддерживаемых математических функций и операторов с наглядными примерами передачи аргументов.
• Многофункциональные ядра (Multi-function kernels): Главная функция main не является обязательной. Программа в терминах MIL может иметь произвольное имя, причем в одном файле допускается объявлять сразу несколько функций. Это дает потенциальную возможность снизить накладные расходы на этапе компиляции.

Фильтруем полезное

Проверять все 163 файла вручную — напрасная трата времени. Я решил автоматизировать данный этап. Используя наработки прошлых тестов, я доработал обертку на языке Rust под названием ane-bridge и создал вспомогательную утилиту ane-compile. Она принимает на вход путь к .mil-файлу, самостоятельно ищет файлы весов (blobfile), раскладывает их по временным директориям и автоматически патчит относительные пути вида /../ в тексте программы.

Запустив утилиту пакетно через find -exec, я смог быстро отфильтровать рабочие конфигурации от тех, которые ANE исполнять отказывается.

Светлое фильтрованное

Автоматическое тестирование сразу вскрыло ряд жестких ограничений платформы ANE. Компилятор принципиально не поддерживает следующие инструкции:

• Операторы scatter* и slice_update: собрать тензор из оригинального массива и обновленного участка привычным образом не получится.
• Инструкции write_state и read_state: прямое управление состоянием (stateful-модели) на ANE не работает.
• Диапазонные размеры RangeDims: сопряжение с произвольными диапазонами отклоняется компилятором (работают только фиксированные перечисления EnumeratedShapes).

Эти ограничения поставили под вопрос простую реализацию обновлений KV-кэша для языковых моделей, ведь без стандартного обновления значений в памяти задачу нужно было решать нестандартным путем.

База

Писать синтаксис MIL вручную в текстовом редакторе чрезвычайно неудобно: компилятор от Apple придирчив к малейшим мелочам, автодополнения нет, а поиск причин сбоя сборки отнимает много времени. Чтобы решить эту проблему, я написал собственный инструмент ProgramBuilder на Rust.

Я остановился на линейной модели сборки: операции дописываются последовательно в конец программы, а к ранее созданным тензорам можно обращаться по их именам. Входы задаются явно, а финализация выходов происходит итоговым вызовом метода done(), после чего сформированный код передается в Kernel::compile.

Программы, индексы и метаданные

MIL-программа может включать несколько исполняемых функций со своими входами и выходами. Однако ранние тесты показали, что компилятор ANE меняет порядок объявленных буферов по собственному усмотрению — сортирует их по алфавиту, размеру или другим внутренним критериям.

Как выяснилось, у загруженного объекта ANEModel можно запросить структуру modelAttributes. Она содержит исчерпывающие данные: номера программ, идентификаторы буферов, соответствующие им имена, типы и точные размеры. Написав парсер для этих метаданных, я избавился от хардкода и получил возможность корректно передавать и считывать буферы исключительно по их именам.

Загружаем веса

Главный нюанс нейромодуля Apple — аппаратные вычисления выполняются строго в формате fp16. Входные данные можно передавать в fp32 или int32, но чип всё равно конвертирует их внутри себя. При этом популярный формат bf16, в котором распространяется большинство LLM на HuggingFace, не поддерживается вовсе.

Для быстрой работы с типами данных в Rust я задействовал библиотеку half. Чтобы ускорить загрузку весов с диска, я применил конвертацию из bf16 в fp16 «на лету» с использованием многопоточности через rayon и считыванием файлов через memmap2 и safetensors. Это полностью исключило лишние аллокации памяти.

Первый блин

Для стартового эксперимента я выбрал модель Qwen3-0.6B — она компактная, быстрая и обладает простой архитектурой. Сначала я написал референсную реализацию для CPU, убедившись, что модель корректно генерирует текст. Это стал мой базовый ориентир.

При переносе вычислений на ANE проявилась проблема: на первом же слое значения «поплыли». Виновником оказался блок rms_norm. Обычно нормализацию считают в fp32 с помощью принудительного приведения типов, но на ANE этот финт недоступен. Проблема решилась математической перегруппировкой формулы: вынеся часть коэффициентов за скобки, мне удалось добиться корректных результатов в fp16 с точностью совпадения до mae=1e-5 по сравнению с CPU.

Вторая подстава возникла на этапе эмбеддингов: компилятор ANE отказывался исполнять операцию gather. Чтобы разобраться в причинах, я изучил проект ANETools и выставил флаг DebugMask для вызова ANECCompile. Попытка собирать бинарные файлы напрямую без _ANEClient провалилась — macOS отказывалась загружать неподписанный бинарник из-за механизмов codesign. Однако дебаг-логи через ANECCompile дали точный ответ: Unsupported tensor data type: int32.

Как оказалось, тип int32 в ANE подходить только для вычисления размеров и форм, но не для индексов выборки. А в тип uint16 индексы токенов большого словаря Qwen банально не помещаются. Поэтому операцию эмбеддингов (как и сэмплирование) пришлось временно оставить на CPU, а всю середину сети — умножение матриц matmul, нормализацию rms_norm и нарезку тензоров — полностью забрать на ANE.

Key-Value Cache

Для организации KV-кэша я разбил программу внутри каждого слоя пополам — между проекцией QKV и блоком Attention. При переходе к вычислению ровно одного токена за шаг проявилось еще одно аппаратное требование ANE: последняя размерность тензора в памяти должна иметь шаг (stride), кратный 32. Если логический размер тензора равен [256, 1], в памяти он обязан занимать [256, 32].

Поскольку прямого обновления памяти через slice_update или scatter на ANE нет, я применил хитрый трюк:

1. С помощью gather собирается тензор обновлений полной длины контекста.
2. С помощью оператора select и маски типа tensor<bool> происходит выбор между старым значением кэша и новым токеном.

Этот ход позволил обновлять контекстное окно полностью на нейрочипе без гонок данных и без участия CPU.

It’s alive

Практический запуск показал важную особенность ANE: цепочка из четырех мелких вызовов ядер выполняется быстрее, чем один объединенный монолитный блок (fused mega-kernel). Попытка свести все методы модели в единый вызов компилятора лишь замедляла инференс.

В итоговой схеме получился следующий пайплайн: эмбеддинг на CPU, 28 слоев Qwen3-0.6B на ANE (по 4 вызова на слой) и сэмплирование на CPU. Однако 112 системных вызовов ANE на каждый токен создавали огромные накладные расходы. Профилирование в утилите Instruments показало, что больше половины времени чип простаивает в ожидании межпроцессных ответов XPC между моим приложением и системным демоном Apple.

Chaining — больше не опция

Попытки задействовать приватный класс _ANEChainingRequest завели в тупик — в самой операционной системе macOS он нигде не используется, а его вызовы приводили к падению драйвера. Зато решение обнаружилось прямо в стандартном объекте ANERequest, а именно в параметре completionHandler.

Вдрив асинхронную отправку запросов с синхронизацией через atomic-wait, я заставил ANE принимать задачи непрерывным потоком. В результате накладные расходы сократились в 3–4 раза, график использования ANE в Instruments стал плотным, а нагрузка на CPU существенно снизилась.

Что дальше?

Модель Qwen3-0.6B подтвердила работоспособность подхода. Более крупные модели (4B и 8B) также успешно запускаются на ANE, хотя скорость упирается в пропускную способность памяти при лимите энергопотребления всего в 3 Вт.

В моделях семейства Qwen3.5 три из четырех слоев используют архитектуру Gated DeltaNet. Это дает константное время обработки при любом размере контекста, но требует адаптации специфических операций (в частности, depthwise convolution), над чем я сейчас и работаю.

Следующим шагом станет квантизация для запуска моделей масштаба 14–32B, а также упаковка созданного Rust-фреймворка в локальный OpenAI-совместимый API сервер, к которому можно будет подключить любого автономного агента или редактор кода.

Дай потыкать

Все наработки по реверс-инжинирингу ANE, включая обертку на Rust и линейный билдер программ, открыты в моем репозитории на GitHub. Пулы с тестами, бенчмарками и оптимизациями приветствуются!

(Небольшое редакционное замечание: в процессе публикации прошлых исследований некоторые коммерческие сервисы вроде Durev VPN без разрешения использовали мои оригинальные тексты для создания рекламных видеороликов на YouTube. Обращайте внимание на первоисточники!).

Источник: habr.com