IBM Research и UC Berkeley адаптировали эволюционный фреймворк K-Search для работы с Apple MLX. Результат: на ядре Attention автоматически сгенерированный код достигает 0.97x от нативного MLX-ядра, а на SSM-ядре Mamba получено 20x ускорение префилла за счёт добавления параллельного редукционного скана, отсутствующего в стандартной библиотеке mlx-lm. Ключевой вывод исследования - узким местом генерации высокопроизводительного кода становится не способность LLM писать код, а качество аппаратного контекста, который подаётся на вход.
Разработчики, работающие с Apple Silicon, знают главную боль: прямые порты CUDA-ядер на Metal дают удручающую производительность. Причина не в слабости железа, а в архитектурных различиях, которые наивный перенос просто игнорирует. K-Search решает эту проблему системно - через слой перевода CUDA-примитивов в архитектурно-корректные Metal-эквиваленты и эвристический поиск, который оперирует не синтаксисом, а аппаратными ограничениями.
Почему перенос CUDA на Apple Silicon - это вызов
GPU Apple и NVIDIA построены на разных философиях. NVIDIA использует дискретную архитектуру с выделенной высокоскоростной памятью и большим объёмом разделяемой памяти на потоковый мультипроцессор - до 164 KB на A100. Apple Silicon применяет unified memory architecture, где GPU и CPU делят общую память, а разделяемая память ограничена 32 KB на потоковую группу. Это фундаментальное различие ломает наивные порты.
Ширина SIMD-групп тоже разная: warp в CUDA - 32 потока, SIMD-группа в Metal - 32 потока, но модель исполнения отличается. Аппаратные инструкции вроде exp2, которые на NVIDIA выполняются за один такт, на Apple GPU требуют другого подхода к разложению. Прямой перенос кода, написанного под warp-синхронное программирование CUDA, натыкается на барьеры потоковых групп Metal и даёт падение производительности в 3-5 раз.
Мы уже видели похожие проблемы в других проектах. Ручные Metal-ядра на Mojo для GPT-2 124M показали, что даже тщательно написанный код на bf16 работает лишь в 1.1x быстрее fp32, тогда как MLX использует тензорные ядра для 2x ускорения. Разрыв между «работает» и «работает быстро» определяется именно качеством адаптации под конкретную архитектуру.
K-Search и MLX: как эволюционный поиск заменяет ручную оптимизацию
K-Search - эволюционный фреймворк, который ищет оптимальную реализацию GPU-ядра через мутации и отбор. Вместо того чтобы инженер неделями перебирал варианты тайлинга, развёртки циклов и размещения данных в памяти, алгоритм делает это автоматически. Исходная версия K-Search была заточена под CUDA и генерировала ядра, конкурирующие с ручными оптимизациями. Адаптация под MLX добавляет принципиально новый компонент: слой перевода, который преобразует не код, а саму постановку задачи оптимизации.
Слой перевода: от CUDA-примитивов к архитектурно-корректным Metal-эквивалентам
Стандартный подход к портированию - взять CUDA-код и переписать синтаксис на Metal. K-Search делает иначе. Слой перевода разбирает CUDA-примитивы на семантические блоки: операции над разделяемой памятью, барьеры синхронизации, паттерны доступа к глобальной памяти, использование warp-уровневых операций. Каждый блок маппится на Metal-эквивалент с учётом аппаратных ограничений целевого GPU.
Пример: когда CUDA-ядро использует 64 KB разделяемой памяти на блок, слой перевода не пытается уместить это в 32 KB Apple GPU. Вместо этого он передаёт эвристическому поиску ограничение «максимум 32 KB разделяемой памяти на потоковую группу» и позволяет алгоритму найти альтернативную стратегию - другой тайлинг, другую раскладку данных, использование регистрового файла для компенсации. Эволюционный поиск получает не синтаксическое дерево, а пространство допустимых решений с чёткими границами.
Аппаратная инструкция exp2, которая на CUDA выполняется аппаратно, на Metal раскладывается в последовательность операций. Слой перевода передаёт поиску эту информацию как стоимость операции, и алгоритм может принять решение: использовать exp2 с известной задержкой или заменить аппроксимацией, если это даст общий выигрыш за счёт лучшего перекрытия вычислений и обращений к памяти.
MLX выступает целевой платформой, предоставляя доступ к низкоуровневым примитивам Metal через Python-интерфейс. Это позволяет K-Search генерировать код, который напрямую использует возможности чипа без прослоек вроде MPSGraph, где часть оптимизаций скрыта за абстракцией. Скрытые возможности Apple M5, такие как поддержка INT8-активаций, остаются неиспользованными именно потому, что высокоуровневые фреймворки не передают аппаратный контекст нижележащим оптимизаторам. K-Search снимает это ограничение.
Результаты: 0.97x на Attention и 20x на префилле Mamba
Исследователи протестировали подход на двух принципиально разных ядрах: Attention - вычислительно-интенсивное ядро с большим объёмом операций на байт данных, и SSM-ядро Mamba - ядро с рекуррентной структурой, где узким местом является последовательная обработка. Результаты показывают, что метод работает в обоих режимах.
Ядро Attention: почти нативная скорость без ручного тюнинга
На ядре Attention эволюционный поиск с полным контекстом из флагманских реализаций достиг 0.97x от нативного MLX-ядра. Это означает, что автоматически сгенерированный код уступает тщательно оптимизированной ручной реализации всего 3%. Условия тестирования: стандартный multi-head attention, размер батча 1, длина последовательности 2048 токенов, размер эмбеддинга 4096, 32 головы. Замеры проводились на MacBook Pro с M3 Max.
Отставание в 3% объясняется не ошибками генерации, а фундаментальным ограничением: нативные MLX-ядра используют ручную настройку под конкретные размеры тензоров, тогда как K-Search генерирует параметризованный код, работающий для диапазона размеров. Это компромисс между универсальностью и пиковой производительностью, и 3% - разумная плата за автоматизацию.
| Реализация | Latency (ms) | Throughput (tok/s) | Относительная производительность |
|---|---|---|---|
| Нативное MLX-ядро | 12.4 | 165.2 | 1.00x |
| K-Search + MLX | 12.8 | 160.0 | 0.97x |
| Наивный порт CUDA | 38.7 | 52.9 | 0.31x |
SSM-ядро Mamba: 20x ускорение префилла благодаря параллельному редукционному скану
С ядром Mamba ситуация интереснее. Префилл в SSM-моделях - это этап предзаполнения, когда модель обрабатывает входную последовательность перед началом генерации. Стандартная реализация в mlx-lm выполняет префилл последовательно: каждый токен обрабатывается за шагом, и время растёт линейно с длиной последовательности. Для последовательности в 32 768 токенов это занимало 840 мс.
K-Search в процессе эволюционного поиска обнаружил возможность добавить параллельный редукционный скан - алгоритм, который преобразует последовательную операцию в дерево параллельных редукций. Этот метод отсутствовал в mlx-lm, потому что стандартная реализация Mamba ориентировалась на инференс с генерацией по одному токену, где префилл не был узким местом. Но для длинных последовательностей параллельный скан даёт радикальный выигрыш: 42 мс вместо 840 мс, ускорение в 20 раз.
Параллельный редукционный скан работает так: вместо обработки токенов один за другим, алгоритм строит бинарное дерево, где на каждом уровне выполняются независимые операции. Для последовательности длиной N требуется log₂(N) шагов вместо N. На GPU с его массовым параллелизмом это превращает операцию из latency-bound в throughput-bound. Похожие оптимизации в llama.cpp - чанкованный SSD matmul для Mamba-2 - дают до 22% ускорения префилла на NVIDIA GPU, но 20x - это результат именно автоматического поиска, который нашёл неочевидную алгоритмическую замену.
| Длина последовательности | mlx-lm (ms) | K-Search + MLX (ms) | Ускорение |
|---|---|---|---|
| 2048 | 52 | 8.7 | 6.0x |
| 8192 | 210 | 17.5 | 12.0x |
| 32768 | 840 | 42 | 20.0x |
Аппаратный контекст - новое узкое место генерации кода
Главный вывод исследования переворачивает привычное представление о возможностях LLM в генерации кода. Способность модели писать синтаксически корректный код на CUDA или Metal - не проблема. Современные LLM справляются с этим уверенно. Проблема в том, что без точных данных об аппаратных ограничениях модель генерирует код, который работает, но не использует возможности железа.
LLM, получившая запрос «напиши эффективное ядро attention для Apple M3», не знает, что разделяемая память ограничена 32 KB, что ширина SIMD-группы - 32, что latency обращения к глобальной памяти составляет порядка 300 тактов, а к разделяемой - 20 тактов. Модель может выдать синтаксически правильный код с разумным тайлингом, но тайлинг будет оптимизирован под абстрактное GPU, а не под конкретный чип. Результат - производительность на уровне 30-40% от нативной.
K-Search решает эту проблему инверсией подхода. Вместо того чтобы просить LLM сгенерировать код, фреймворк передаёт эвристическому поиску аппаратные ограничения как параметры пространства решений. LLM в этом пайплайне может использоваться для генерации начальных популяций ядер или для предложения мутаций, но ключевое - она получает на вход не абстрактное «напиши ядро», а спецификацию: «доступно 32 KB shared memory, 256 регистров на поток, latency global load 300 cycles, целевой размер тензора M×N×K». Качество этого контекста напрямую определяет качество результата.
Параллель с промпт-инжинирингом прямая: спецификация железа - это системный промпт для генерации кода. Чем точнее и полнее спецификация, тем ближе результат к оптимальному. 0.97x на Attention и 20x на префилле Mamba - прямое следствие качества аппаратного контекста, переданного поиску.
Как применить подход в своих проектах
K-Search с MLX-бэкендом доступен в репозитории проекта. Для воспроизведения результатов на Apple Silicon потребуется macOS 14.0+, чип M1 или новее, Python 3.10+ и установленный MLX. Адаптированный слой перевода CUDA-в-Metal пока работает с ограниченным набором примитивов, но покрывает основные паттерны: операции matmul, attention, свёртки и редукции.
Быстрый старт с K-Search на Apple Silicon
Минимальный сценарий для оптимизации пользовательского ядра:
pip install mlx ksearch # Клонирование репозитория с адаптированным бэкендом git clone https://github.com/ibm/ksearch-mlx cd ksearch-mlx # Запуск эволюционного поиска для attention-ядра python -m ksearch.optimize \ --kernel attention \ --backend mlx \ --target M3Max \ --population 256 \ --generations 100 \ --output optimized_attention.metal
Фреймворк принимает описание ядра на предметно-ориентированном языке или в виде референсной CUDA-реализации, автоматически применяет слой перевода и запускает эволюционный поиск. Результат - файл с Metal-кодом, готовый к интеграции в проект.
Ограничения подхода: наибольший выигрыш достигается на вычислительно-интенсивных ядрах и ядрах с последовательными зависимостями, где пространство оптимизаций велико. Для простых поэлементных операций ручная реализация часто достаточна. Требования к железу: поиск по популяции из 256 индивидов на 100 поколениях занимает 20-40 минут на M3 Max, для M1 время увеличивается до 1.5-2 часов. Альтернативные подходы к инференсу на Apple Silicon через WebGPU показывают, что аппаратно-специфичные оптимизации критичны для любого бэкенда - разница между универсальным и специализированным ядром измеряется кратными приростами производительности.
Что дальше: эволюция автоматической оптимизации ядер
Результаты K-Search на Apple Silicon - часть более широкого тренда: автоматическая генерация высокопроизводительного кода переходит от синтаксического уровня к семантическому. Следующий шаг - расширение слоя перевода на другие бэкенды: AMD ROCm, Intel oneAPI, Qualcomm Adreno. Каждый новый бэкенд требует не переписывания кода, а добавления спецификации аппаратных ограничений в слой перевода.
Аппаратно-осведомлённые промпты становятся стандартным инструментом. Ожидается, что будущие версии фреймворков инференса будут включать формализованные описания железа, которые автоматически подаются на вход оптимизаторам и LLM. Это снимет текущее ограничение, когда разработчик должен вручную исследовать характеристики чипа перед написанием эффективного ядра.
Текущие ограничения подхода: слой перевода покрывает не все CUDA-примитивы, а эволюционный поиск требует значительных вычислительных ресурсов на этапе оптимизации. Однако для production-систем, где ядро оптимизируется один раз и используется миллионы раз, эти затраты окупаются. Исследование IBM и UC Berkeley задаёт вектор: будущее не за ручным написанием ядер под каждую архитектуру, а за системами, которые принимают на вход алгоритм и аппаратную спецификацию, а на выходе дают близкий к оптимальному код.