ИИтог

K-Search перенёс опыт CUDA-ядер на Apple Silicon: Mamba ускорилась в 20 раз

исследованияРедакция ИИтогBerkeley AI Research🔥66Рейтинг ИИтог66Влияние68Новизна62Релевантность85Неожиданность48Участники52Достоверность85
K-Search перенёс опыт CUDA-ядер на Apple Silicon: Mamba ускорилась в 20 раз

Команда UC Berkeley (первый автор, Шии Цао, Berkeley Sky Lab, публикация в блоге Berkeley AI Research, BAIR) расширила K-Search, эволюционный фреймворк для автоматической оптимизации GPU-ядер, новым бэкендом для MLX, фреймворка машинного обучения Apple для чипов Apple Silicon. Проблема, которую решает работа: экосистема CUDA за десятилетия накопила глубокую, вручную настроенную экспертизу для внимания, моделей пространства состояний (SSM) и других критичных операций, а более молодые экосистемы вроде Apple Silicon этой глубины лишены. MLX запускает модели корректно, но оставляет много производительности на столе: paged attention, оптимизированные ядра сканирования SSM и fused MoE-роутинг либо отсутствуют, либо реализованы наивно без учёта архитектуры.

K-Search работает как итеративный поиск: получив наивное ядро и спецификацию железа, языковая модель в роли «инженера по производительности GPU» предлагает следующую оптимизацию, модель для генерации кода реализует кандидата, кандидат компилируется и замеряется на реальном железе, а результат возвращается в поиск. Состояние поиска хранится как «модель мира», дерево решений, где каждый узел получает оценку (общий рейтинг от 0 до 10, уверенность от 0 до 1, влияние на пропускную способность памяти, давление на регистры и соответствие железу); дерево растёт от раунда к раунду, а при застое поиск переключается на другую ветку. В экспериментах авторов обе роли, рассуждение и написание кода, выполняла одна модель, Gemini 3.5 Pro Preview. Оригинальная статья про K-Search (Цао и соавторы, 2026) показала, что фреймворк превосходит OpenEvolve и ShinkaEvolve на CUDA-ядрах из библиотеки FlashInfer (GQA decode, MLA decode, MLA prefill, MoE) при одинаковом бюджете в 120 итераций.

Ключевой вклад этой работы, слой перевода CUDA-экспертизы в термины MLX/Metal. Он включает таблицы соответствий примитивов (например, CUDA shared соответствует threadgroup-памяти Metal, но с лимитом в 32 КБ против 48 КБ у Nvidia; warp_reduce лучше заменять на MMA; __syncthreads() становится threadgroup_barrier(); пропускная способность H100 около 3,35 ТБ/с HBM3 против примерно 400 ГБ/с общей памяти у M3 Max, разница, которая меняет набор выгодных оптимизаций), специфичные для MLX паттерны (регистровые редукции через simd_shuffle_xor, «трюк с exp2», замена exp(x) на 2^(x·log₂e), чтобы использовать быструю аппаратную инструкцию exp2 на Apple Silicon вместо обычной экспоненты) и переиспользуемые утверждения, свойства, которые должно сохранять эталонное ядро, а не код для копирования.

Результаты проверялись в двух сценариях. Для ядра внимания сравнили три конфигурации: наивный baseline, эволюцию без дополнительного контекста и полный слой перевода. Полный контекст поднял производительность с 0,26x до 0,97x от скорости нативного, «эталонного» ядра внимания MLX от Apple, самостоятельно переоткрыв ключевые приёмы FlashAttention-2: тайлинг через threadgroup-память, онлайн-softmax, транспонирование K для доступа к памяти и трюк с exp2. Для ядра SSM модели Mamba (тест на mamba-370m в f16 на M1 Max 64 ГБ) decode-производительность составила 152 токена в секунду против 116 токенов в секунду у ядра из открытой реализации mlx-lm, а ускорение prefill достигло до 20 раз относительно того же mlx-lm; третьим ориентиром служила эталонная PyTorch-реализация mamba.py.

Ключевые факты

  • Berkeley AI Research (BAIR) расширила эволюционный фреймворк K-Search (Цао и соавторы, Berkeley Sky Lab) бэкендом для MLX, фреймворка Apple для Apple Silicon
  • Слой перевода сопоставляет примитивы CUDA и Metal/MLX с учётом реальных ограничений железа: например, лимит threadgroup-памяти 32 КБ у Apple против 48 КБ у Nvidia, и разница в пропускной способности H100 (~3,35 ТБ/с) против M3 Max (~400 ГБ/с)
  • На ядре внимания полный слой перевода поднял скорость с 0,26x до 0,97x от нативного MLX-ядра Apple, самостоятельно переоткрыв приёмы FlashAttention-2
  • На ядре SSM для Mamba (mamba-370m, M1 Max 64 ГБ) decode вырос до 152 токенов в секунду против 116 у community-реализации mlx-lm, а prefill ускорился до 20 раз
  • Обе роли в поиске, рассуждение об оптимизациях и написание кода ядер, выполняла одна модель, Gemini 3.5 Pro Preview

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

Работа показывает, что десятилетия ручной оптимизации CUDA-ядер можно частично перенести на новую архитектуру автоматически, а не переоткрывать с нуля. Это меняет расклад для более молодых экосистем вроде Apple Silicon: MLX уже умеет запускать модели корректно, но без глубокой аппаратной настройки (paged attention, оптимизированные SSM-сканы, fused MoE) теряет производительность. K-Search с переводным слоем впервые показал, что этот разрыв можно закрыть эволюционным поиском поверх CUDA-знаний, а не рукописной портацией.

Кому это важно

В первую очередь, инженерам, которые запускают локальный инференс языковых моделей на Mac (unified-память Apple Silicon особенно удобна для моделей на 7, 70 млрд параметров), и разработчикам самого MLX. Дальше, исследователям систем и компиляторов для ИИ, которые ищут способы переносить экспертизу между аппаратными экосистемами (кастомные ускорители, не только Apple), и командам, которые строят инструменты автоматической оптимизации ядер.

Как это применить

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

Можно ли доверять

Источник, официальный блог Berkeley AI Research (BAIR), результат работы Berkeley Sky Lab; текст ссылается на конкретную предшествующую статью (Цао и соавторы, 2026) с воспроизведёнными графиками сравнения против OpenEvolve и ShinkaEvolve на CUDA-ядрах FlashInfer. Это повышает доверие к методологии. При этом сами цифры по MLX/Mamba, самоотчёт авторов в блоге, а не результат независимой рецензируемой публикации или стороннего повторения экспериментов.

Риски и подводные камни

Без переводного слоя результат составлял всего 0,26x от эталонной скорости (источник не уточняет, идёт ли речь о наивном baseline или об эволюции без дополнительного контекста), то есть без заранее подготовленных человеком таблиц соответствий и паттернов автоматика сама не находит нужные оптимизации. Это означает, что подход не полностью автоматический: под каждый новый тип операции или новую аппаратную экосистему придётся заново готовить экспертный переводной слой. Benchmarks проведены на одной модели (Gemini 3.5 Pro Preview) и ограниченном наборе железа (M1 Max, M3 Max), как результаты обобщаются на другие модели и чипы Apple, из текста не ясно; таблица с полным сравнением decode-производительности против эталонной PyTorch-реализации mamba.py в источнике обрывается до итоговых цифр.