Продолжаем человеческий 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.
