29 июля 2026 года стало датой, отмечающей новый этап в развитии аппаратного и программного обеспечения для ИИ. Рынок стремительно меняется: появляются чипы с разной архитектурой, а инструменты кодирования ИИ теперь справляются с задачами, на которые раньше требовались месяцы, за считанные минуты.

Ключевую роль в успехах ИИ играют GPU-ядра. Однако перенос ядра с оборудования одного производителя на другое — сложная задача, требующая поиска новых оптимизаций. Экосистема CUDA за десятилетия накопила богатый опыт в создании ядер — от оптимизированных механизмов внимания до моделей пространства состояний. При этом новые экосистемы, такие как Apple Silicon, хотя и развиваются быстро, пока не могут похвастаться таким же уровнем экспертизы.

Чтобы решить эту проблему, был разработан подход K-Search. Он базируется на эволюционной системе поиска ядер, которую ввели Cao et al. в Berkeley Sky Lab. Суть подхода в том, чтобы использовать существующие CUDA-ядра как базу знаний и адаптировать их под Apple Silicon, а не создавать ядра с нуля. K-Search позволяет достичь производительности, близкой к уровню эксперта на Apple Silicon: например, ускорение составляет 0,97 раза по сравнению с родным MLX Attention-ядром, а предварительная загрузка ускоряется до 20 раз по сравнению с реализацией mlx-lm в сообществе для ядра Mamba SSM.

Фреймворк MLX от Apple активно внедряется с конца 2023 года. Он позволяет выполнять локальный вывод ИИ без затрат на облако. Особенно привлекателен MLX для моделей среднего размера (7B–70B параметров) на чипах серии M благодаря архитектуре с единой памятью. Однако в MLX отсутствует ряд критически важных для производительности ядер, которые в экосистеме NVIDIA считаются стандартными. Среди них — страничное внимание, оптимизированные ядра сканирования SSM, объединённая маршрутизация MoE.

K-Search — это эволюционная система оптимизации ядер, созданная Shiyi Cao в UC Berkeley Sky Lab. Её принцип работы следующий: при заданном простом ядре и спецификации оборудования запускается итеративный цикл оптимизации. LLM определяет, какие оптимизации попробовать дальше, модель написания кода генерирует варианты ядер, варианты компилируются и тестируются на реальном оборудовании. Результаты измерений возвращаются в поиск, который продолжает совершенствоваться — развивает перспективные направления и отбрасывает тупиковые, пока производительность не сойдётся.

Алгоритм 1: K-Search via co-evolving world models — поиск чередует выбор наиболее перспективного действия, создание и оценку кода до стагнации улучшения и эволюцию модели мира через операции вставки, обновления и обрезки (из Cao et al. (2026)). Поиск опирается на документ Spec, который кодирует правила оборудования, шаблоны оптимизации и математические ограничения. Это позволяет избежать генерации некорректного кода и обеспечить эффективную компиляцию и выполнение.

В экспериментах одна модель (Gemini 3.5 Pro Preview) выполняет сразу две роли: поддерживает состояние рассуждений и пишет ядра. Процесс начинается с этапа рассуждений, который имитирует работу инженера по производительности GPU-ядер. Инженер классифицирует ядро (reduction, scan, attention/softmax и др.), переписывает эталонные вычисления в канонической форме, составляет схему размещения данных и шаблонов доступа, выдвигает гипотезы о вероятных узких местах (пропускная способность, задержка, вычисления или синхронизация) в каждом режиме выполнения. Только после этого генерируются кандидаты на оптимизацию — по одному изменению за итерацию.

Модель мира в K-Search представлена не плоским списком вариантов, а деревом решений (префиксным деревом). Каждый путь от корня к листу — это полный план оптимизации, а ветви-братья — конкурирующие альтернативы. Каждый узел оценивается по нескольким критериям: overallrating в диапазоне [0, 10], уверенность в [0, 1], влияние на пропускную способность памяти, нагрузку на регистры и соответствие вычислениям и оборудованию. Дерево сохраняется и растёт между раундами: при уточнении идеи добавляется дочерний узел, а если в течение определённого времени улучшения отсутствуют, поиск переходит к альтернативной ветви. Пример узла модели мира в середине выполнения для attention kernel — в ```
{
"action": "Replace the threadgroup-memory softmax reduction
with a register-only reduction: each SIMD group
owns 8 query rows and reduces across lanes with
simdshufflexor, removing a threadgroupbarrier.",
"difficulty1to5": 4,
"impacts": {
"memorybandwidth": 8,
"registerpressure": 4, // risk: spill if Br > 8
"computehwfit": 9 // SIMD width 32; keep tile 8x8
},
"overallrating0to10": 8,
"confidence0to1": 0.7
}
```.

Рабочий процесс K-Search состоит из трёх фаз:
1. Выбор действия — из фронтира извлекается наиболее перспективный узел действия на основе оценочного приоритета модели мира $V$.
2. Локальное уточнение — стохастическая политика $\pi{\mathrm{code}}$ генерирует конкретные реализации до стагнации.
3. Обновление модели мира — LLM анализирует траекторию и обновляет дерево поиска через вставку (добавление новых действий), обновление ($V$, например, $u{11}$ снижается с 0,9 до 0,6) и обрезку (удаление менее перспективных узлов, например $u_{10}$).

Стратегия поиска оценивалась на CUDA-ядрах из FlashInfer. Для GQA decode, MLA decode, MLA prefill и MoE K-Search показал более стабильное улучшение по сравнению с OpenEvolve и ShinkaEvolve при бюджете в 120 итераций. На трёх запусках K-Search достиг более высоких результатов по лучшим найденным оценкам поиска, производительности ядер для каждой рабочей нагрузки и распределениям ускорения, чем OpenEvolve и ShinkaEvolve на четырёх CUDA-ядрах FlashInfer (рис. 3, воспроизведено из Cao et al. (2026)).

Для переноса K-Search на Apple Silicon был создан нативный бэкенд MLX и реализован адаптер задач, специфичный для MLX. MLX task backend в k_search/tasks/ обрабатывает компиляцию и выполнение ядер на Apple Silicon через Metal/C++ API MLX. Также были обновлены подсказки для генератора ядер для написания и модификации ядер Metal/MLX и интегрирован бенчмаркинг для MLX с использованием утилит измерений mlx.core.

Главная задача при переносе знаний оптимизации из CUDA в MLX — избежать ситуации, когда простые попытки перенести ядро CUDA через LLM дают синтаксически верный, но архитектурно неверный код. Например, могут быть выбраны неправильные размеры тайлов, недопустимые примитивы или сделаны несовпадающие предположения о памяти.

Слой перевода включает:
* таблицы сопоставления концепций: глоссарий примитивов CUDA и их эквивалентов в MLX/Metal с жёсткими ограничениями (например, shared сопоставляется с памятью группы потоков Metal с ограничением в 32 КБ (против 48 КБ у NVIDIA), warpreduce — с MMA, _syncthreads — с threadgroupbarrier(memflags::memtg));
* сопоставление пропускной способности: H100 с ~3,35 ТБ/с HBM3 сопоставляется с M3 Max с ~400 ГБ/с унифицированной DRAM;
* подсказки и паттерны для MLX: паттерны на уровне кода для операций без прямого эквивалента в CUDA (например, сокращения строк на основе регистров с использованием simdshufflexor в компоновке тайла MMA 8×8, «трюк exp2» — замена $exp(x)$ на $exp2(x \log_2 e)$ для ускорения softmax);
* повторно используемые утверждения: поведение экспертных ядер переосмыслено как свойства, которые должен сохранять эволюционный поиск.

Была проведена оценка трёх конфигураций ядра внимания MLX для Apple Silicon:
1. наивный базовый вариант;
2. чистая эволюция без дополнительного контекста;
3. слой перевода с полным контекстом, который предоставляет оптимизатору знания о реализации, специфичные для архитектуры (например, FlashAttention-2).

Конфигурация «Полный контекст» успешно реализует продвинутые стратегии (двойное буферирование, развёртывание циклов) и достигает почти экспертного уровня производительности. Рост с 0,26× до 0,97× скорости современного ядра внимания Apple показывает значимость слоя перевода. С полным контекстом эволюционировавшее ядро самостоятельно находит ключевые оптимизации FlashAttention 2: разбиение памяти группы потоков, онлайн-softmax, транспонирование K для доступа к памяти, «трюк exp2».

Чтобы проверить обобщающую способность K-Search, его применили к ядру модели пространства состояний (SSM), используемому в Mamba. Проведено сравнение эволюционировавшей реализации с реализацией сообщества MLX (mlx-lm) и эталонной реализацией PyTorch (mamba.py) на M1 Max. Оценка проводилась на mamba-370m f16, M1 Max 64GB. В таблице 1 представлены показатели пропускной способности для mamba-370m (f16, M1 Max 64GB):
* decode: mlx-mamba (ours) — 152 tok/s, mlx-lm (community) — 116 tok/s, mamba.py — 40 tok/s;
* prefill L=512: mlx-mamba — 5751 tok/s, mlx-lm — 329 tok/s, mamba.py — 1089 tok/s;
* prefill L=1024: mlx-mamba — 6010 tok/s, mlx-lm — 327 tok/s, mamba.py — 1127 tok/s;
* prefill L=2048: mlx-mamba — 6612 tok/s, mlx-lm — 326 tok/s, mamba.py — 1092 tok/s;
* prefill L=4096: mlx-mamba — 6743 tok/s, mlx-lm — 339 tok/s, mamba.py — 1042 tok/s.

MLX-mamba достигает примерно в 20 раз более высокой пропускной способности prefill по сравнению с mlx-lm за счёт реализации параллельного сканирования для SSM, тогда как mlx-lm обрабатывает токены последовательно. Mamba.py работает медленно на prefill и decode, так как это эталонная реализация PyTorch, которая использует CPU или MPS на Apple Silicon, не применяя оптимизаций, доступных в бэкенде MLX Metal.

Работа ведётся в нескольких направлениях: поддержка новых архитектур (в том числе разработка ядер для IBM Spyre AIU), добавление новых ядер (например, paged attention и fused MoE routing), улучшение интеграции с циклом эволюции K-Search. Работа выполнена IBM Research на основе K-Search из UC Berkeley Sky Lab.

Для воспроизведения результатов необходимо:
1. Клонировать и установить (```
git clone
cd K-Search

uv pip install openai wandb
uv pip install git+
).
2. Установить учётные данные, открыв соответствующий скрипт в scripts/ и задав три переменные вверху (
KSEARCHROOT=/path/to/K-Search
APIKEY=your-llm-api-key
).
3. Запустить поиск ядра (

Optimize Flash Attention on Apple Silicon (world-model mode)

bash scripts/macflashattention_wm.sh

Or a Mamba SSM kernel, e.g. the selective scan

bash scripts/mambaselectivescanfwdwm.sh
```).

Shiyi Cao

Gal Bloch

Assaf Toledo

Michael Factor

Gil Vernik

Joseph E. Gonzalez

K-Search

Cao et al., 2026

@article{cao2026k,
 title={K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model},
 author={Cao, Shiyi and Mao, Ziming and Gonzalez, Joseph E and Stoica, Ion},
 journal={arXiv preprint arXiv:2602.19128},
 year={2026}
}

RSS feed