Перенести быстрое GPU-ядро с железа NVIDIA на Apple Silicon всегда означало одно: найти эксперта, который заново изобретёт все оптимизации с нуля. Исследователи из IBM Research и UC Berkeley показали, что это больше не обязательно. Их система K-Search с новым слоем трансляции CUDA в MLX сама вывела attention-ядро, которое работает со скоростью 0.97x от нативного ядра Apple, и ускорила prefill Mamba SSM примерно в 20 раз против community-реализации mlx-lm.

K-Search это эволюционный фреймворк поиска GPU-ядер, где LLM играет роль инженера по производительности: предлагает оптимизации, генерирует код, компилирует и бенчмаркует его на реальном железе, а затем улучшает результат итерациями. Работа опубликована в блоге BAIR 29 июля 2026 года, авторы из UC Berkeley и IBM Research.

Почему MLX и в чём проблема

Фреймворк MLX от Apple стал стандартом локального инференса на Mac с конца 2023 года. Унифицированная память Apple Silicon отлично подходит для моделей среднего размера, от 7B до 70B параметров на чипах M-серии. Но под капотом зияет дыра: многие критичные для производительности ядра, которые экосистема NVIDIA считает само собой разумеющимися, в MLX либо отсутствуют, либо написаны наивно. Paged attention, оптимизированные scan-ядра для state-space моделей, слитый роутинг для MoE, всё это работает корректно, но оставляет значительную часть производительности железа неиспользованной.

За десятилетия вокруг CUDA накопились тысячи инженерных часов: вручную вылизанные реализации attention, SSM и других критических операций. Вопрос, который поставили авторы: можно ли перенести эту экспертизу автоматически, без команды GPU-специалистов.

Почему существующие подходы не решают проблему

Перенос ядер на новое железо сегодня идёт тремя путями, и у каждого есть предел. Первый путь это ручное переписывание: эксперт по Metal изучает CUDA-ядро и пишет аналог с нуля. Так появились лучшие ядра в MLX, но это месяцы работы редких специалистов, и результат не масштабируется на десятки операторов и каждый новый чип.

Второй путь это высокоуровневые DSL вроде Triton, которые абстрагируют детали архитектуры. Проблема в том, что сам Triton должен быть портирован и настроен под конкретное железо, а его backend для Apple Silicon находится в зачаточном состоянии. Абстракция переносится только тогда, когда под ней уже есть чужая экспертная работа.

Третий путь это наивное использование LLM: дать модели CUDA-код и попросить перевести на Metal. Авторы прямо показывают, почему это ломается. Без глубокого контекста железа модель выдаёт код, который компилируется, но архитектурно неверен: размеры тайлов под 48 КБ shared memory, которых нет в 32 КБ threadgroup-памяти Apple, примитивы warp-редукций, не существующие в SIMD-группах Metal, допущения о пропускной способности, откалиброванные под HBM3 в 3.35 ТБ/с, тогда как M3 Max даёт около 400 ГБ/с. Синтаксис правильный, физика нет.

Как устроен K-Search

На вход система получает наивное ядро и спецификацию железа. Дальше работает итеративный цикл: LLM рассуждает, какую оптимизацию попробовать следующей, кодовая модель генерирует вариант ядра, вариант компилируется и замеряется на железе, а измерения возвращаются в поиск. В их экспериментах обе роли играла одна модель, Gemini 3.5 Pro Preview.

Интересна структура рассуждения. Это не плоский список «что попробовать», а дерево решений, которое авторы называют world model. Каждый путь от корня к листу это полный план оптимизации, а соседние ветви конкурируют между собой. Каждый узел получает оценку: общий рейтинг от 0 до 10, уверенность от 0 до 1 и прогноз влияния на пропускную способность памяти, давление на регистры и соответствие вычислительному блоку. Если лучший результат не улучшается несколько раундов, поиск откатывается и исследует альтернативную ветвь.

Типичный узел выглядит как конкретное действие: «заменить softmax-редукцию в threadgroup-памяти на редукцию только в регистрах: каждая SIMD-группа владеет 8 строками запросов и редуцирует по линиям через simd_shuffle_xor, убирая threadgroup_barrier». К такому действию прилагается оценка сложности (4 из 5), влияние на память (8), риск регистрового спилла (4) и соответствие железу (9, потому что ширина SIMD равна 32 и тайл 8x8 ложится идеально).

В оригинальной статье K-Search (Cao et al., 2026, arXiv:2602.19128) эта стратегия обошла OpenEvolve и ShinkaEvolve на ядрах FlashInfer при одинаковом бюджете в 120 итераций, на задачах GQA decode, MLA decode, MLA prefill и MoE.

Слой трансляции CUDA в MLX

Главная находка работы не в том, что K-Search запустили на MLX, а в том, как передали знания. Просто скормить LLM CUDA-ядро с просьбой «портируй» не работает: модель выдаёт синтаксически валидный, но архитектурно неправильный код, с неверными размерами тайлов, несуществующими примитивами и ошибочными допущениями о памяти.

Слой трансляции состоит из трёх частей. Первая это таблицы соответствия концепций: структурированный глоссарий CUDA-примитивов и их эквивалентов в MLX/Metal с жёсткими ограничениями. Например, __shared__ маппится в threadgroup-память Metal, но с лимитом 32 КБ против 48 КБ у NVIDIA. __syncthreads() становится threadgroup_barrier(mem_flags::mem_tg). А главное: ~3.35 ТБ/с HBM3 у H100 против ~400 ГБ/с унифицированной DRAM у M3 Max. Такая разница в пропускной способности переворачивает представление о том, какие оптимизации вообще стоит делать.

Вторая часть это паттерны уровня кода для операций без прямого аналога в CUDA. Например, редукции строк в регистрах через simd_shuffle_xor в тайловой раскладке 8x8 MMA, или «трюк exp2»: замена каждой экспоненты в softmax на экспоненту по основанию 2 по формуле e^x = 2^(x·log2 e). Это математически точно и позволяет использовать быструю аппаратную инструкцию fast::exp2() Apple вместо платной конверсии основания в рантайме.

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

Attention: от 0.26x до 0.97x от ядра Apple

Авторы сравнили три конфигурации MLX-ядра attention для Apple Silicon: наивный бейзлайн, чистая эволюция без дополнительного контекста и полный контекст со слоем трансляции, где оптимизатор получает знания из высокопроизводительных ядер вроде FlashAttention-2. Такая постановка позволяет изолировать вклад именно слоя трансляции.

Результат говорит сам за себя: скачок с 0.26x до 0.97x от скорости state-of-the-art ядра Apple. С полным контекстом эволюционировавшее ядро независимо переоткрыло ключевые оптимизации FlashAttention-2: тайлинг через threadgroup-память, online softmax, транспозицию K для доступа к памяти и тот самый трюк exp2. Система не копировала код FlashAttention, а пришла к тем же решениям через поиск, направляемый структурированным знанием.

Mamba SSM: 20x быстрее на prefill

Чтобы проверить генерализацию за пределами attention, авторы применили K-Search к ядру state-space модели Mamba. Здесь узкое место другое: не softmax, а рекуррентное обновление состояния. Сравнение шло против community-реализации mlx-lm и эталонной mamba.py на M1 Max с 64 ГБ, модель mamba-370m в f16.

Числа из таблицы впечатляют. Decode: 152 токена в секунду у их mlx-mamba против 116 у mlx-lm и 40 у mamba.py. Prefill при длине 512: 5751 ток/с против 329 у mlx-lm и 1089 у mamba.py. При L=1024: 6010 против 327 и 1127. При L=2048: 6612 против 326 и 1092. При L=4096: 6743 против 339 и 1042.

Откуда берётся двадцатикратная разница на prefill? Всё сводится к одному: mlx-lm не реализует параллельный scan для SSM. Рекуррентность h_t = a_t·h_{t-1} + b_t выглядит по своей сути последовательной, но каждый шаг можно записать как пару (a_t, b_t) с ассоциативной операцией комбинирования. Ассоциативность означает, что всю последовательность можно вычислить параллельным prefix-scan за O(log N) зависимых шагов вместо O(N). mlx-lm обрабатывает токены по одному, оставляя большую часть вычислительной мощности Apple Silicon без работы. Эволюционировавшее Metal-ядро применяет scan и загружает GPU по-настоящему.

Показательно, что выигрыш проявляется именно на prefill, где вся последовательность доступна для параллельного сканирования, а на decode разницы почти нет: там приходит один токен за шаг, и параллелить нечего. А mamba.py медленная на обоих режимах, потому что это эталонная PyTorch-реализация, которая на Apple Silicon откатывается на CPU или MPS без аппаратно-специфичных оптимизаций Metal.

Стоит разобрать, что означают эти числа практически. Prefill это фаза обработки промпта, и она доминирует в сценариях с длинным контекстом: суммаризация документов, анализ кода, RAG с большими выдержками. При prefill в 329 ток/с обработка контекста в 4096 токенов на mlx-lm занимает больше двенадцати секунд, а при 6743 ток/с меньше секунды. Для интерактивного локального ассистента это разница между «заметно тупит» и «отвечает мгновенно». И обратите внимание на асимптотику: у mlx-mamba пропускная способность растёт с длиной последовательности, от 5751 до 6743 ток/с, потому что параллельный scan амортизирует накладные расходы, а у mlx-lm она стоит на месте около 330 ток/с.

Почему это важно

Главный вывод авторов переворачивает привычное представление об автогенерации кода. Узким местом оказалась не способность LLM писать Metal-код, а качество контекста и ограничений, которые ей дали. Слой трансляции превращает накопленную NVIDIA-экспертизу в действенные указания для Apple Silicon, а эволюционный поиск делает остальное.

Это сигнал для всей индустрии гетерогенного железа. Новые ускорители появляются быстрее, чем сообщество успевает накопить для них ядерную экспертизу. Если метод переносится, а авторы подчёркивают, что он не специфичен для MLX, то разрыв между «железо вышло» и «софт его использует на полную» может сократиться с лет до недель. Авторы уже работают над ядрами для IBM Spyre AIU, paged attention и слитым MoE-роутингом.

Часто задаваемые вопросы

Что такое K-Search простыми словами?

K-Search это система, где LLM автоматически оптимизирует GPU-ядра эволюционным поиском: предлагает изменение, генерирует код, измеряет скорость на реальном железе и улучшает результат итерациями. Дерево решений с оценками направляет поиск в перспективные ветви и отсекает тупиковые.

При чём здесь CUDA, если ядра пишутся для Apple?

Экспертные CUDA-ядра кодируют десятилетия знаний об оптимизации. Слой трансляции преобразует эти знания в таблицы соответствия, паттерны и ограничения для Metal/MLX, чтобы LLM рассуждала о стратегиях, а не копировала синтаксис, который не подходит чужой архитектуре.

Можно ли воспроизвести результаты?

Да, MLX-бэкенд построен поверх открытого репозитория K-Search (github.com/caoshiyi/K-Search). Нужно клонировать репозиторий, установить зависимости через uv, прописать API-ключ LLM в скрипте и запустить готовые сценарии для Flash Attention или Mamba selective scan на Mac.

Итог

K-Search со слоем трансляции CUDA в MLX показал: автоматический перенос ядерной экспертизы между архитектурами работает. Attention-ядро достигло 0.97x от нативного Apple, Mamba SSM получила 20x на prefill против community-реализации, и всё это без команды GPU-экспертов. Если у вас есть Mac на M-чипе, попробуйте воспроизвести результаты из открытого репозитория: это лучший способ понять, на что способен эволюционный поиск ядер.