Продолжаем человеческий ANE reverse engineering. В этот раз разберёмся, чем таким интересным можно накормить компилятор, чтобы получить желаемый результат - Qwen, запущенный без GPU и CPU.

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

Реверсить бинарник компилятора целиком не хотелось. Примеры, доступные в сети, довольно однообразны, и все копируют друг друга - все проекты, так или иначе, ссылаются на исследование Maderix Claude. Где же ещё искать примеры под чип Apple, если не в системе Apple?

find /System/Library -name ‘*.mil’ -type f

И вот у нас в руках 163 (у вас может получиться другое число) файла с примерами MIL-кода. Валидного. Правда, с оговоркой - валидный он как CoreML MIL (на который, кстати, есть какой-никакой референс). Не все из этих файлов удастся скомпилировать и запустить под ANE.

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

• Dynamic shapes: в первую очередь я полез искать что-то вроде ? - классическое обозначение "я не знаю какой shape тут будет". И нашёл - работает это через аннотацию FlexibleShapeInformation. Аннотация содержит DefaultShapes и RangeDims/EnumeratedShapes - вернёмся сюда чуть позже.

• Новые типы данных: как "языковые" (tuple, list, dict), так и для данных (pixel_buffer, tensor_buffer, state<T>).

• Operation set: конечно, много операций. Много того, что можно сделать с вашими данными. С примерами валидных аргументов, конечно же.

• Multi-function kernels: оказывается, main - условность. Функция (программа, давайте называть её так, дальше это сделает вещи проще) может называться как угодно. И их может быть несколько в одном файле. Уменьшаем overhead компиляции ещё сильнее (но это не точно).

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

Проверять все 163 файла руками - зачем тратить 2 часа на то, что можно автоматизировать за неделю 10 минут? У нас уже есть (с предыдущей статьи) обёртка на rust, которая берёт MIL и собирает/запускает его. Правда, тут же возникают нюансы:

1. Во многих моделях, помимо MIL, лежат файлы с весами (blobfile), которые надо аккуратно разложить во временные директории перед сборкой.

2. В некоторых MIL эти файлы ещё и находятся на уровень выше, из-за чего пути к этим файлам содержат /../ - надо пропатчить MIL и убрать эту историю.

В общем, за 10 минут не уложился, но всё таки написал example к ane-bridge - который так и назвал, ane-compile, потом ещё пригодится. Путь к mil на входе, с остальным разберётся сам. Вызываем через find -exec, и идём наливать кофе.

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

Для начала, чего точно нет ни в одном файле, собранном под ANE:

• scatter*, slice_update - собрать тензор из оригинального и обновлённой части обычным способом не получится.

• write_state, read_state - stateful-модели в ANE не запускаем.

• RangeDims - диапазонные динамические размеры не поддерживаются (правда, есть EnumeratedShapes).

Идея запустить LLM на ANE без CPU начала казаться сложнее чем в начале - как же обновлять KV-кэш, если нельзя записать новые значения в их место?

База

- Товарищ полковник, машина не заводится!
- Фигня, поехали, потом заведешь!

Писать MIL руками - дело неблагородное. Даже зная, что можно писать. Компилятор придирчивый, подсветки нет, autocompletion нет, блокнот - не наш путь. Разбираться, почему отвалилась компиляция очередного шедевра инженерной мысли - нет уж, спасибо. Строим Builder. Поначалу, была мысль построить полноценный графовый, но остановился на линейном - просто очередная операция дописывается в конец программы, предыдущие можно референсить по названию. Входы задаём явно, выходы указываем завершающим вызовом Builder. Получается почти красиво:

let program = ProgramBuilder::new()
  .build_info(HashMap::from([("ane_fn_name".to_owned(), "qwen_out".to_owned())]))

  .func("qwen_out_b1")
  .input("w_model_norm", TensorType::new(TensorDtype::Fp16, &[hidden_size])?)
  .input("w_lm_head", TensorType::new(TensorDtype::Fp16, &[self.config.vocab_size, hidden_size])?)
  .input("i_hidden", TensorType::new(TensorDtype::Fp16, &[hidden_size])?)
  .rms_norm("i_x_norm", "i_hidden", "w_model_norm", rms_eps, 0, nwp1)?
  .matmul("logits", "i_x_norm", "w_lm_head", false, true)?
  .done(["logits"])

  .done();

Kernel::compile("qwen_out", &self.bridge, &program.mil(), &[])?

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

Программ (именно так в терминологии Apple называются выполняемые куски кода внутри MIL) может быть несколько. У каждой - свои входы, свои выходы. Мы их как-то обозначили, дали названия, типы, и отдали компилятору. А потом нам надо запустить конкретную программу, дав ей все необходимые буферы входов и выходов.

Ранние эксперименты показали, что компилятор не особо хочет сохранять порядок того, что мы наобъявляли. Что-то сортирует по алфавиту, что-то по размеру, что-то просто как хочет. Но ведь так не бывает - должна быть какая-то система? И да, она есть. Если у загруженной ANEModel спросить modelAttributes - она вернёт очень подробную структуру, в которой будут указаны все номера программ, все айдишники буферов, соответствующие им названия, размеры, и кучу других метаданных. Так что убираем хардкод и пишем честный парсер - с этого момента можем по названию передавать и получать что угодно, без регистрации и смс.

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

Тут было бы всё просто, если бы не одно но: ANE считает только в fp16. На вход ему можно скормить много чего - fp32, int32, да даже int8, но тогда надо конвертировать данные в самом ANE. А формат bf16, в котором распространяются почти все модели на huggingface, вообще не поддерживается. Хорошо, что есть крейт half - который позволяет удобно работать с fp16/bf16 прямо в Rust. Ещё одной прослойкой меньше - в предыдущих работах данные преимущественно ходили в fp32, с постоянными cast туда-обратно. Однако, для загрузки весов, готового метода bf16-to-fp16 нет - не беда, берём спецификацию Qwen и просим написать сниппет. Упрощаем валидацию (nan/inf в весах нас не интересуют), подключаем rayon, и грузим тензоры с диска потоком (memmap2, safetensors, ноль лишних аллокаций).

Первый блин

Начать я решил с Qwen3-0.6B - весит немного, работает шустро, математика простая. Референс на numr был написан за час - с переписыванием кода между языками у меня редко возникали проблемы. Модель запустилась на CPU, и уверенно начала писать ответы. Baseline был пройден - наивная реализация, без кэша, без оптимизаций. Пора было переписывать на ANE - блок за блоком. Не без приключений - ну как же иначе.

Я уже говорил, что ANE считает только в fp16? С первого же слоя данные "поплыли" - и виновником оказался rms_norm. Обычно его считают в fp32 (принудительным cast), но нам эта опция недоступна. Благо, тот же Qwen быстро придумал вспомнил трюк, который всё исправил: ввести scale, что-то вынести за скобки, что-то перегруппировать, и в итоге получить математически корректный результат в переменных низкой размерности. Прогнал тесты между CPU fp32 и преобразованным вариантом на ANE - точность совпала до mae=1e-5 (max=1e-4). Модель начала отвечать более вразумительно.

Вторая подстава возникла с embedding: компилятор ANE отказывался давать мне gather. Пора было ковырнуть, могу ли я достать более адекватные сообщения об ошибках - и нашёл проект ANETools. В нём вызов ANECCompile реализован напрямую, и вынесено большое количество настроек - интересовал меня DebugMask.

В этот момент я подумал - а почему бы не собирать бинарники напрямую, а не через _ANEClient? А вот потому что. Скомпилированная модель не загружалась - и в Console (стандартная утилита macOS, собирающая вообще все логи со всех уголков системы) нашлись упоминания codesign. Тут пришло понимание, что загружать неподписанные бинарники в ANE просто так не получится. Чем их подписывать? А чёрт его знает, но ANEClient через XPC вызывает внутренние механизмы Apple, и оно работает, а как говорится, работает - не трогай. Тем не менее, ANECCompile позволял увидеть отладочную информацию - а она нужна ровно в тот момент, когда что-то пошло не так. Именно туда (в обработку ошибки компиляции) я и добавил вызов in-process сборки - с дебагом, с логами, с удобствами.

А с логами пришло и понимание: Unsupported tensor data type: int32. Как покажут дальнейшие исследования, int32 можно использовать только для shape/size, и в около-константных операциях типа concat (для тех же shape). Но для индексов он не подходит. В меньший тип данных (uint16) не влезут индексы токенов - словарь сильно больше. Поэтому embedding остался на CPU (пока).

Ладно, отстанем от "одноразовых" операций (sampling я тоже оставил на CPU, так как воевать с argmax/topk желания не было) - что там в середине? А середина ушла на ANE целиком - matmul и rms_norm делали своё дело, slice_by_size честно резал тензоры, concat склеивал обратно. Attention я сразу делал на цепочке matmul - быстрый тест подтвердил показания Claude о том, что attn_mask игнорируется. В целом - всё работало. Пора было идти в рейд на финального босса.

Key-Value Cache

Окей, делаем шаг назад. Нарезаем наш program прямо посередине слоя - между QKV projection и собственно Attention. Временные буферы проблемой не являются. Проверяем - всё ещё работает.

Ограничиваем вычисления: теперь мы считаем не все токены, а ровно один. Переписываем input, все остальные размеры пересчитываются на лету (спасибо builder!), проверяем - опа, упало. Что упало? Размер буфера не совпал. Так узнаём, что последний dimension тензора имеет stride кратный 32. Для входов и выходов это важно - если мы говорим, что тензор [256, 1] - в памяти он должен быть как [256, 32]. Расточительно по байтам и неудобно читать/писать данные с CPU, поэтому фиксируем это в своей биологической памяти и докидываем reshape где надо.

Теперь нам надо забирать "новые" K/V из первой программы, писать в наш "кэш" (который был нашим временным буфером между ядрами), и запускать оставшуюся часть. Окей, всё ещё работает, но не прикольно - как убрать из этой цепочки CPU?

Наивное решение - slice_update. У нас есть оригинал (старый кэш), новые данные, и диапазон который надо заменить - всё сходится? Нет, эта операция не поддерживается. Cannot support standalone slice_update - что бы ни значил этот standalone. Декомпиляция ANECompiler показала, что этой ошибкой завершается любая попытка парсинга этого оператора. Примеры из системы показывают, что этой операцией надо пользоваться в связке со state<T>, read_state, write_state - но данные операции тоже не компилируются в ANE.

Чуть более замороченно - scatter. Есть оригинал, есть update, индексы собрать не проблема - и снова мимо, Unsupported MIL operation “scatter”.

Спустя несколько минут медитации на список операций из документации, в голову приходит безумная идея. У нас есть оригинал. У нас есть select, позволяющий выбирать между A и B в зависимости от tensor<bool, [...]> cond. И у нас есть gather, позволяющий набрать данных из другого тензора. Набираем из "обновлений" по индексам, чтобы получить тензор "полной" длины - вне обновляемого диапазона "забираем" нулевой индекс, в целом нам всё равно что там будет. Обновляемый диапазон маскируем, и делаем select между полным (старым) кэшем, и новыми данными. Эффективно? Ну, выглядит так себе, но надеемся на внутренний оптимизатор. Будет ли работать? Сейчас узнаем... да!

It's alive

Критический инсайт по результатам: в попытке собрать данную цепочку вызовов обратно, выяснилось, что 4 мелких вызова отрабатывают быстрее, чем один fused mega-kernel. То же самое справедливо для multi-program kernel: попытка свести все методы одной модели в один вызов компилятора замедлила inference в разы. Вишенкой на торте стал механизм enumerated shapes: он, по факту, разворачивается на этапе компилятора в несколько независимых программ, и наследует проблему замедления от multi-program. То-есть - "можно, а зачем но не нужно".

В остальном - у меня был пайплайн из CPU-embed, full-ANE-layers (28 слоёв для Qwen3-0.6B), по 4 вызова на слой, и CPU-sampler в конце. В сочетании с builder, это дало возможность без переписывания всего сделать chunked-prefill, ускорив TTFT. Всё было хорошо... Кроме ~112 вызовов ANE на каждый токен. Ну, то-есть, overhead. По результатам профилирования через Instruments (заботливо показывающего использование ANE), было выяснено, что больше половины времени не работает никто. ANE закончил предыдущий вызов, через XPC моему процессу летит ответ "готово", тот возвращается в поток, где я запускаю новый вызов, и тот через XPC летит обратно в сервис, который занимается ANE, и дальше в ядро. Долго. Медленно. Не нравится.

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

Изначально я хотел посвятить этой теме отдельную статью (с разносом того, куда опять понесло ЫЫ), но потом понял, что в этом нет никакого смысла - тема достаточно короткая. Всё как в прошлой статье - нейронка увидела знакомое слово, подумала что "вот оно", и пошла ковырять то, что на самом деле не нужно.

_ANEChainingRequest. Очевидно, что он занимается "связкой" нескольких вызовов вместе. Но правда ли он тут нужен? Возможно. Но исследования в этой области завели нейронку в тупик. Очевидно, почему - её in-memory path был фундаментально несовместим с тем, что на вход ожидают методы вокруг chaining. Замена на обычную ANEModel (которая уже была сделана в предыдущей серии) дала толчок вперёд, реверс функций показал, что надо положить в остальные параметры, а подробные логи дали понимание порядка вызовов. Через некоторое время все методы возвращали "ок"... Но код не работал. Системные логи рассказывали о каких-то необработанных событиях и assert в низкоуровневом драйвере.

В этот момент я решил поискать вызовы этого метода в системе. Кто-то (CoreML? CoreAI?) должен был работать с ANE эффективно. Но... нет. Использований этого класса в системе не нашлось (как и в случае с ANEInMemoryModel). Зато нашлось кое-что другое. Ответ всё это время был у меня в руках, и был этим ответом... completionHandler. Это обычная практика для "универсальных" методов - хочешь синхронный - просто вызывай, хочешь асинхронный - дай callback и жди пока его позовут.

Такой callback был у самого обычного ANERequest - того самого, который использовался с самого первого исследования. Задание ему "какого-то" блока сразу сломало всё - код стал полностью асинхронным, и даже повесил macOS - пришлось перезагружать ноут принудительно. Вставив в обработчик мьютекс для синхронизации (в дальнейшем заменённый на atomic-wait), и добавив вызов sync, я получил желаемый результат - теперь работа для ANE отправлялась в очередь подряд, а в месте, где мне нужен был её результат, CPU вставал ждать. График использования ANE в Instruments стал плотнее, а использование CPU - заметно меньше. Можно ли улучшить этот результат? Естественно, сейчас overhead составляет около 10-15%, но это в 3-4 раза лучше, чем было в синхронном варианте. Но для этого надо решить проблему с "нативным" chaining - если она, конечно, решается.

Что дальше?

Qwen3-0.6B - это, конечно, здорово. Но хотелось что-то поумнее. 4B и 8B тоже запускаются и работают - правда, медленно. Всё-таки, 16 гигабайт матриц в 3W питания ворочать тяжело.

Семейство Qwen3.5 - это уже сильно лучше. Кроме повышения качества ответов (модель "новее" и предсказуемо "умнее" при тех же размерах), там в 3 из 4 слоях задействованы более эффективные Gated DeltaNet - что означает рекуррентную обработку, константное время и память при любом контексте. Но тут меня ждали новые подставы от ANE: depthwise convolution из коробки не заработал. Новые трюки для эффективной обработки - WIP.

Следующий уровень - квантизация. Если 8B ещё использовать комфортно (16GB RAM + контекст), то модели уровня 14-32B уже в адекватное количество RAM не влезают. А хочется Qwen3.8-27B, ещё и с MTP, и со всем фаршем... Ну и MoE, куда же без Qwen3.6-35B-A3B.

Превращение данного фреймворка в универсальный комбайн пока под сомнением. Training? Конечно, можно пробросить binding в python, где посчитать градиенты и оптимизаторы (всё, конечно же, асинхронно и без CPU). И, скорее всего, оно будет работать. Но мой интерес в другой области.

Обёртка данного фреймворка в OpenAI-совместимый API - вот то, куда я точно буду это развивать. Подключение к любому клиенту, будь то редактор кода или автономный агент.

Дай потыкать

После обширного рефакторинга, я всё же готов показать код: GitHub. Наработки по Qwen3.5, OpenAI API, и квантизации будут появляться по мере стабилизации. Pull Requests с тестами, бенчмарками, оптимизациями и другими полезностями - you are welcome.