Продолжаем человеческий 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 и собирает/запускает его. Правда, тут же возникают нюансы:
-
Во многих моделях, помимо MIL, лежат файлы с весами (
blobfile), которые надо аккуратно разложить во временные директории перед сборкой. -
В некоторых 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.
Автор: DjPhoeniX
