Kernel Forge: LLM-агент ускорил CUDA-ядра до 2,83 раза с помощью MCTS

Большая часть времени работы ИИ-моделей уходит на небольшой набор вычислительных ядер, умножение матриц, свёртку, нормализацию. Их оптимизация, прямой способ снизить задержку и затраты, но традиционно это требует инженеров, вручную пишущих низкоуровневый код для GPU. Агентные системы на базе LLM уже умеют генерировать и оптимизировать такие ядра с меньшим участием человека, но у существующих инструментов авторы отмечают ряд ограничений: их проверяют на случайно сгенерированных тензорах и изолированных ядрах, они выдают отдельный CUDA-код, который разработчику приходится вручную встраивать обратно в проект, работают в основном только с LLM-моделями на PyTorch и слабо поддерживают инспекцию и отладку результатов.

Авторы представили Kernel Forge, открытую сквозную (end-to-end) агентную «обвязку» (harness), которая принимает на вход любую немодифицированную PyTorch-модель. Kernel Forge поддерживает vision-, diffusion- и LLM-нагрузки, а вместо одной линейной цепочки последовательных правок использует Monte Carlo Tree Search (MCTS), метод перебирает сразу несколько путей оптимизации параллельно. В комплекте, графический интерфейс для мониторинга прогресса, просмотра ядер-кандидатов и отладки сбоев.

Систему протестировали на четырёх PyTorch-моделях из областей vision, diffusion и LLM на GPU NVIDIA GB10 (DGX Spark). Всего за 50 итераций оптимизации на ядро Kernel Forge улучшил 14 ядер так, что они стали работать быстрее режима PyTorch eager mode: ускорение в 1,52 раза на adaptive_avgpool2d в ResNet-50, в 1,70 раза на group_norm в Stable Diffusion 3.5 Medium, в 2,83 раза на softmax в Gemma 4 E2B и в 1,54 раза на softmax в Qwen 3.5 35B-A3B.

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

Kernel Forge, открытая агентная «обвязка» на LLM для генерации и оптимизации CUDA-ядер; принимает немодифицированную PyTorch-модель без предварительной подготовки.
Вместо линейной цепочки правок использует Monte Carlo Tree Search (MCTS), перебирает несколько путей оптимизации параллельно, а не по одному кандидату за раз.
Работает с vision-, diffusion- и LLM-нагрузками; включает графический интерфейс для мониторинга прогресса и отладки ядер-кандидатов.
На четырёх моделях (ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B, Qwen 3.5 35B-A3B) на GPU NVIDIA GB10 всего за 50 итераций на ядро оптимизировал 14 ядер быстрее PyTorch eager mode, ускорение от 1,52 до 2,83 раза.
Максимальное измеренное ускорение, 2,83 раза на softmax в Gemma 4 E2B.

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

Большинство времени выполнения ИИ-моделей уходит на узкий набор вычислительных ядер, умножение матриц, свёртку, нормализацию, softmax. Их оптимизация, самый прямой рычаг снижения задержки и затрат на инференс, но раньше для этого нужны были инженеры, вручную пишущие низкоуровневый CUDA-код. Kernel Forge показывает, что агент на базе LLM с поиском по дереву решений (MCTS) может взять на себя эту работу для целой немодифицированной модели, а не только для отдельного изолированного ядра.

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

Инженерам инфраструктуры и MLOps-командам, которые обслуживают модели на GPU и ищут способ снизить стоимость и задержку инференса без ручной оптимизации кода. Полезно и разработчикам vision-, diffusion- и LLM-систем на PyTorch, именно эти три класса нагрузок покрывает инструмент.

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

Kernel Forge, открытый исходный код, который принимает на вход обычную, ничем не изменённую PyTorch-модель и сам ищет более быстрые реализации её узких мест. Графический интерфейс позволяет следить за прогрессом оптимизации, сравнивать ядра-кандидаты и разбирать причины сбоев, не копаясь в логах вручную.

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

Результаты получены на реальном железе (GPU NVIDIA GB10 в составе DGX Spark) и на четырёх реальных моделях из разных областей, ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B, Qwen 3.5 35B-A3B, а не на синтетических тензорах. Авторы указывают конкретные ускорения по каждому ядру, что даёт возможность проверки и воспроизведения.

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

Проверка охвата ограничена: всего четыре модели и 50 итераций оптимизации на ядро, результаты получены на одном типе GPU (NVIDIA GB10) и могут не переноситься напрямую на другое железо. Из общего числа ядер в моделях быстрее эталона стали только 14, то есть метод оптимизирует не каждое ядро подряд, и часть работы всё ещё не даёт выигрыша.