Запуск Qwen на Apple Neural Engine: reverse engineering MIL и компиляция без GPU

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

Для этого необходимо было найти примеры MIL-кода — промежуточного представления, которое использует компилятор CoreML. Автор не стал разбирать бинарник компилятора целиком, а обратился к системным файлам Apple. Команда find /System/Library -name '*.mil' -type f вернула 163 файла с примерами MIL-кода. Эти файлы валидны для CoreML MIL, но не все из них можно скомпилировать и запустить под ANE. Анализ показал, что в них используются динамические формы через аннотацию FlexibleShapeInformation, новые типы данных (tuple, list, dict, pixel_buffer, tensor_buffer, state<T>), а также multi-function kernels — программы могут называться как угодно и их может быть несколько в одном файле.

Однако при фильтрации выяснилось, что ANE не поддерживает ряд операций: scatter, slice_update, write_state, read_state и RangeDims. Это означает, что нельзя обновлять KV-кэш в языковых моделях стандартными методами. Тем не менее, автор нашёл обходной путь: он построил линейный Builder на Rust, который позволяет создавать MIL-программы без ручного написания кода. Пример программы для Qwen включает операции rms_norm и matmul, что достаточно для вычисления логитов.

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

Таким образом, автор продемонстрировал, что запуск LLM на ANE возможен, но требует глубокого понимания внутренностей компилятора и ограничений архитектуры. Этот опыт полезен для разработчиков, которые хотят оптимизировать инференс моделей на устройствах Apple, особенно в сценариях, где важна энергоэффективность и автономность.