Hazy Research: агенты делают CUDA-абстракции ненужными

Стэнфордская исследовательская группа Hazy Research, авторы CUDA-библиотеки ThunderKittens для написания GPU-кернелов, опубликовала пост о том, как их подход к написанию «мегакернелей» изменился за год благодаря ИИ-агентам.

Год назад команда писала мегакернель для моделей Llama. Мегакернель, это единый GPU-кернел, который объединяет несколько вычислительных шагов, чтобы не терять время на запуск отдельных кернелов и синхронизацию между ними; писать его вручную тяжело: нужны сложные структуры данных, синхронизация между потоками, потоковыми блоками (SM) и всей GPU, глубоко вложенная логика. Держать это в голове было невозможно, поэтому команда сделала то, что инженерия делает уже 70 лет, построила слой абстракции (собственно ThunderKittens). Но даже с этим слоем несколько месяцев ушло на борьбу с гонками состояний (race conditions) и взаимными блокировками (deadlock), прежде чем мегакернель для Llama заработал быстро.

В этом году команда построила мегакернель для MoE-модели (mixture-of-experts, «смесь экспертов»), но уже без слоя C++-абстракции. Вместо этого ИИ-агентам напрямую поручили работать со сложностью и писать код, оптимизированный под конкретное железо, с нуля. Авторы отмечают закономерность последнего квартала: задачи, которые и раньше решались без абстракций (например, написание оптимизированного GEMM-кернеля, кернеля матричного умножения), теперь при правильном промпте почти полностью автоматизируются агентом, хотя выбор PTX-инструкций и дизайн warp-специализации команда пока задаёт сама. А вот задачи, для которых раньше требовалась абстракция, написание мегакернеля целиком, агентами пока не автоматизированы: одним запросом мегакернель сегодня не получить. Но агент позволяет держать «абстракцию» прямо в промпте, в неполном и неаккуратном виде, а не в тщательно спроектированных C++-шаблонах: сложность, которая год назад была неподъёмной, стала управляемой, потому что появился «компилятор», способный превращать расплывчатые инструкции в код.

Отсюда, центральный тезис поста: по индукции CUDA DSL как класс, включая любимый авторами ThunderKittens, следующие в очереди на «пенсию», вероятно в следующем году или раньше; насколько глубоко пойдёт этот процесс, авторы не знают. Их аргумент: кодовая база даёт точность, одна и та же машина исполнит её одинаково дважды, но эта точность стоит хрупкости: код привязан к языку, фреймворку, железу и соглашениям, понятным только написавшей его команде. Промпт, противоположность: он расплывчат, зато переносим между разными исполнителями, человеком или машиной, если тот исполнитель достаточно умён, чтобы верно заполнить пробелы. DSL и фреймворки, это способ заранее специфицировать все пробелы для «глупого» исполнителя; если исполнитель (агент) перестаёт быть глупым, у DSL пропадает почва под ногами. Авторы проводят параллель с компилятором: индустрия уже согласилась не писать ассемблер вручную и доверять непрозрачному преобразованию исходного кода, потому что компилятор оказывался прав чаще людей; вопрос сейчас, различие ли это в степени или в роде, и авторы считают, что не в роде.

При этом пост оговаривает условия, при которых слой абстракции вообще можно «уволить». Во-первых, абстракция, это ещё и общая поверхность, на которой держатся переиспользование и код-ревью команд; без неё десять команд, пишущих кернели самостоятельно, получают десять несвязанных наборов проблем вместо одного общего. Во-вторых, слой можно убрать только тогда, когда есть «оракул», который переживёт этот слой, эталонные реализации, допуски по точности вычислений, интуиция о том, как должен выглядеть профилированный график исполнения; там, где такого оракула ни у кого нет, вывод не работает и каркас нужно оставить. В-третьих, авторы прямо называют себя единственной пристрастной выборкой: они убрали абстракцию в области, которую знают очень глубоко, и не берутся утверждать, работает ли то же самое для новичка, который впервые сталкивается с warp-специализацией, абстракции ведь служат ещё и для передачи знаний тем, кто ими пока не владеет.

Вывод поста: смысл кодовой базы смещается, команда собирается и дальше поддерживать, обновлять и любить ThunderKittens (готовятся кернели под будущие GPU Nvidia Vera Rubin), но при этом ожидает, что конкретная реализация станет одноразовым «кэшем» одной компиляции, а не источником истины: доверие и внимание при ревью переносятся со среза кода (diff) на спецификацию и на «оракул», который её проверяет.

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

• Год назад Hazy Research потратила несколько месяцев на борьбу с гонками состояний и deadlock, чтобы вручную построенная C++-абстракция (ThunderKittens) выдала быстрый мегакернель для Llama.
• В этом году команда построила мегакернель для MoE-модели вообще без слоя C++-абстракции, код, оптимизированный под конкретное железо, писали напрямую ИИ-агенты.
• Задачи, которые и раньше решались без абстракции (например, GEMM-кернель), сейчас при правильном промпте почти полностью автоматизируются агентом; написать мегакернель целиком одним запросом агенты пока не могут.
• Вывод авторов: CUDA DSL как класс, включая ThunderKittens, движутся к «пенсии», вероятно в следующем году или раньше, но точный горизонт неизвестен.
• Условие для отказа от абстракции, наличие «оракула» (тестов, эталонных реализаций, допусков по точности), который переживёт этот слой; без такого оракула каркас нужно сохранять.

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

Пост фиксирует смену парадигмы в низкоуровневой GPU-инженерии: раньше сложность мегакернелей приручали тщательно спроектированным C++-слоем абстракции, теперь тот же результат получают, отдав нечёткие инструкции агенту прямо в промпте. Если верен главный тезис, «пока исполнитель был глупым, нужно было специфицировать каждый пробел заранее; когда исполнитель поумнел, у специфицированного слоя пропадает почва под ногами», под вопросом оказывается смысл целого класса инструментов: DSL и фреймворков, которые существуют именно как заранее прописанная спецификация для не умеющего рассуждать компилятора или разработчика.

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

В первую очередь, инженерам, которые пишут и поддерживают низкоуровневые GPU-кернели и CUDA-библиотеки (в том числе создателям и пользователям самой ThunderKittens), а также командам, разрабатывающим DSL и фреймворки для машинного обучения: именно их область работы автор поста прямо называет кандидатом на «пенсию». Шире это касается всех, кто проектирует агентные инструменты для написания низкоуровневого, требовательного к железу кода.

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

Авторы формулируют явное условие, а не общий призыв отказаться от абстракций: слой можно убрать только там, где есть «оракул», эталонная реализация, допуски по точности, независимый способ проверить результат агента, который переживёт удаляемый слой. Там, где такого оракула ни у кого нет, вывод не действует и каркас абстракции нужно сохранить. При этом даже отказавшись от C++-слоя, команда сознательно оставляет себе: цель, инварианты, тесты и накопленное знание о том, что делает кернель корректным на конкретном железе, то есть отказ от библиотеки не означает отказа от знания, которое в ней жило.

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

Источник, сама команда-разработчик ThunderKittens, которая описывает собственный опыт от первого лица, без привлечения независимых данных: в посте нет ни одной цифры производительности, сравнивающей мегакернель на абстракции и мегакернель, написанный агентом, не названы ни конкретный агент/модель, ни точная дата работы, ни имена конкретных авторов внутри команды. Сами авторы прямо оговаривают, что представляют «единственную пристрастную выборку» и не уверены, распространяется ли их вывод за пределы области, которую они знают очень глубоко.

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

Без общего слоя абстракции пропадает разделяемая поверхность для код-ревью и переиспользования: десять команд, каждая генерирующая свой мегакернель агентом, получают десять несвязанных наборов проблем вместо одного общего набора, который раньше решался и проверялся один раз для всех. Отдельная неопределённость, доступность знания для новичков: абстракции традиционно служат ещё и способом передать накопленный опыт тем, кто в теме не разбирается, и авторы прямо признают, что не знают, работает ли агентный подход так же хорошо для человека, впервые встретившего warp-специализацию, или удобство «без абстракций» доступно только уже обученным специалистам.

«Доверие смещается на уровень выше. Мы вглядываемся в спецификацию и в оракул, а не в код-ревью. Реализация становится одноразовой, кэшем одной конкретной компиляции, а не источником истины.»

— Hazy Research, блог Stanford