Если вы когда-нибудь писали CUDA-код под GPU-ядра для трансформеров, вы знаете эту боль: чтобы выжать максимум из H100 или Blackwell, приходится считать байты shared memory, следить за warp-синхронизацией и вручную раскладывать tensor core instructions. PyTorch скрывает эту сложность, но забирает контроль. Triton даёт больше свободы, но всё ещё абстрагирует hardware. А raw CUDA — это сила, но цена — сотни строк бойлерплейта на каждое ядро. Stanford Hazy Research Lab предлагает третий путь: ThunderKittens — компактный DSL, встроенный прямо в CUDA, который даёт ощущение PyTorch-примитивов, но сохраняет доступ к железу.

Что такое ThunderKittens

ThunderKittens — это embedded DSL внутри CUDA C++ от лаборатории Hazy Research в Стэнфорде. Его центральная идея проста: предоставить ограниченный, но выразительный набор абстракций, которые напрямую отображаются на иерархию GPU — от глобальной памяти до регистров warp'а и tensor cores. Вместо того чтобы писать сотни строк ручного управления memory layout и warp-координации, разработчик работает с тайлами, векторами и pipeline-шаблонами, которые под капотом генерируют оптимизированный машинный код.

Как пишут сами авторы: «Несмотря на очевидную потребность в бесчисленных техниках для задействования всех этих аппаратных возможностей, наш центральный технический вывод состоит в том, что для многих ИИ-ядер действительно существует небольшое число ключевых абстракций, которые могут упростить процесс написания high-performance kernels». Это не попытка заменить CUDA — это попытка найти золотую середину между продуктивностью и производительностью.

Почему это важно для scaling laws

Tri Dao изtogether.ai формулирует scaling laws через простое уравнение: Intelligence per Dollar = (Intelligence per FLOPS) × (FLOPS per Dollar). Первый множитель — это алгоритмы и данные, второй — эффективность железа. Исследователи постоянно придумывают новые архитектуры attention, новые схемы обучения MoE и новые техники квантизации. Но любой алгоритмический прорыв обесценивается, если он не укладывается в реальные GPU с их ограничениями по памяти, пропускной способности и compute throughput.

Современные модели зависят от кастомных ядер: FlashAttention, grouped GEMM для MoE, различные варианты квантизированных матричных умножений. Каждое такое ядро — это ручная оптимизация под конкретное железо. Чем быстрее исследователи могут прототипировать и отлаживать эти ядра, тем быстрее движется весь фронт ИИ-разработки. ThunderKittens решает именно эту задачу — сокращает время от идеи до работающего optimized kernel с минимальной потерей производительности.

Тайловые абстракции: строительные блоки

В основе ThunderKittens лежит идея тайла — прямоугольного фрагмента матрицы, который отображается на иерархию памяти GPU. Базовый тайл TK имеет фиксированную высоту 16 строк, а ширина зависит от типа данных: 16 колонок для bf16/fp16 и 32 колонки для fp8. Это не случайный выбор — он напрямую следует из того, как tensor core instructions экспонируют свои фрагменты на Hopper и Blackwell.

На Hopper инструкции wgmma работают с формами вроде m64n256k16, где M-сторона зафиксирована на 64, а N может масштабироваться. TK представляет такие операции как композицию базовых 16×16 фрагментов: тайл st_bf<64,64> — это логически 4×4 решётка из 16×16 кусков. Для fp8 базовый фрагмент расширяется до 16×32, и та же логика применяется с поправкой на ширину. Таким образом, разработчик думает в терминах «тайл 64×64», а не «64 варпа, каждый со своим фрагментом, синхронизировать через barriers».

TK определяет несколько ключевых типов тайлов. Global layout (gl) описывает тензор в глобальной памяти GPU и инкапсулирует TMA-метаданные для асинхронных копий. Shared tile (st) представляет данные в shared memory с правильным swizzling — специальным переупорядочиванием элементов, которое минимизирует bank conflicts при доступе warp'ов. Register tile (rt) отображается на регистры потоков и напрямую участвует в tensor core операциях. Векторные типы sv и rv дополняют картину для warp-level и thread-level вычислений соответственно.

Compute и memory movement

После того как данные описаны через тайлы, над ними нужно что-то вычислять. TK предоставляет два семейства операций: maps и reductions. Map сохраняет форму тайла, применяя скалярную операцию поэлементно — например, сложение двух тайлов даёт тайл того же размера. Reduction сворачивает форму: вектор в скаляр, тайл в строку или столбец. Различие принципиально: map — это elementwise broadcast, reduction — это агрегация.

Для tensor cores TK предоставляет отдельный слой абстракций. На Hopper warpgroup из четырёх warp'ов (128 потоков) кооперируется на одной MMA-операции через инструкции wgmma. На Blackwell пятая генерация tensor cores использует инструкции tcgen05, которые работают с ещё более крупными формами. TK скрывает различия между поколениям: разработчик пишет rt_fl<64,64> = rt_fl<64,64> @ rt_fl<64,64> + rt_fl<64,64>, а фреймворк сам эмитирует правильные wgmma или tcgen05 в зависимости от целевой архитектуры.

Перед вычислениями данные должны пройти через иерархию памяти. TK предоставляет shared_allocator — bump-аллокатор, который разбивает динамический SMEM-буфер на типизированные регионы. Вместо ручного вычисления смещений и приведения типов разработчик пишет al.allocate, 2>() и получает массив из двух тайлов нужного типа. Это мелочь, которая экономит десятки строк кода и устраняет целый класс багов с неправильным alignment.

Pipeline-шаблоны: от абстракций к целым ядрам

Отдельные тайлы и операции — это строительные блоки. Pipeline-шаблоны — это готовые каркасы для типичных паттернов выполнения. TK определяет три уровня: kernel schema (что ядро читает и пишет), hook interface (какие функции должен реализовать разработчик) и pipeline engine (как эти функции оркестрируются во время выполнения).

Шаблон lcf (load-compute-finish) реализует классический double-buffered pipeline: пока один набор тайлов вычисляется, следующий загружается из глобальной памяти. Разработчик определяет только producer::load для загрузки данных и consumer::compute для их обработки — всё остальное, включая синхронизацию и переключение буферов, берёт на себя движок. Шаблон lcsf добавляет этап finish для редукции результатов или записи финальных тайлов. Если задача не укладывается в готовые шаблоны, разработчик падает на нижний уровень абстракций и пишет ядро вручную — но всё ещё с тайлами, а не с raw pointers и inline PTX.

Стратегии тайлового планирования

Как только задача разбита на тайлы, возникает вопрос: как распределить их по SM? Наивный подход — один тайл на один CTA — страдает от wave quantization: если число тайлов не кратно числу SM, последняя волна CTAs недоиспользует GPU. TK поддерживает persistent scheduling, где один CTA берёт на себя несколько тайлов подряд, устраняя простои между волнами.

На Blackwell NVIDIA добавила Cluster Launch Control — аппаратный механизм, позволяющий кластерам из нескольких CTAs кооперироваться на shared memory уровне. TK экспонирует это через соответствующие pipeline-шаблоны. Ещё одна техника, supergrouping, улучшает L2-reuse за счёт того, что соседние CTAs работают с близкими тайлами, повышая вероятность попадания в кэш второго уровня.

Swizzling и memory layout: почему порядок данных имеет значение

Одна из самых неочевидных, но критически важных деталей GPU-программирования — это то, как данные физически размещаются в shared memory. Когда 32 потока warp'а одновременно обращаются к shared memory, GPU разбивает эти обращения на транзакции по 128 байт. Если все потоки читают последовательные адреса, транзакция одна. Если адреса разбросаны — транзакций может понадобиться 32, и производительность падает в десятки раз. Это явление называется bank conflicts.

ThunderKittens решает эту проблему через автоматический swizzling — переупорядочивание элементов тайла в shared memory так, чтобы соседние потоки warp'а обращались к разным банкам. Для 16-битных типов TK использует 32-байтные swizzle-регионы: восемь 16-битных значений в строке упаковываются в один 16-байтный чанк, а два таких чанка формируют 32-байтный регион. Для fp8 регион удваивается до 64 байт, потому что базовый фрагмент шире. Разработчик не управляет этим вручную — layout тайла инкапсулирует swizzling, и данные всегда размещаются оптимально для coalesced доступа.

Это особенно важно при переходе между поколениями GPU. На Hopper оптимальные формы wgmma имеют фиксированную M-сторону 64 и масштабируемую N-сторону до 256. На Blackwell архитектура tensor cores меняется: микробенчмаркинг Hazy Research показывает, что Blackwell ведёт себя как 128×128 систолический массив. Если запустить ядро с 64×64 выходным тайлом на Blackwell, заполняется только половина M-стороны и половина N-стороны — примерно четверть пиковой производительности. Для полной утилизации нужны тайлы 128×128 или даже 256×128 в режиме 2SM MMA. TK позволяет параметризовать тайловые размеры и перекомпилировать ядро под целевое железо без переписывания логики.

Сравнение с экосистемой: где TK занимает свою нишу

GPU-программирование для ИИ существует на спектре от полной абстракции до полного контроля. PyTorch находится на верхнем конце: исследователь пишет tensor expressions, а фреймворк сам решает, какое ядро запустить. С PyTorch 2 TorchDynamo и TorchInductor генерируют Triton-ядра, которые часто конкурентоспособны с ручными реализациями. Но этот путь терпит неудачу, когда нужна нестандартная memory layout, кастомная синхронизация или экспериментальная инструкция, которую Inductor ещё не поддерживает.

Triton сдвигает баланс в сторону контроля: разработчик явно описывает тайловую структуру, tiling и memory access patterns, но всё ещё не видит warps, shared memory barriers и tensor core instructions. Это хороший уровень для большинства кастомных ядер, но недостаточный для максимальной оптимизации под конкретное железо. CUDA C++ и CUTLASS дают полный контроль, но требуют глубокой экспертизы и порождают огромное количество шаблонного кода. PTX — ассемблер GPU — открывает доступ к каждой инструкции, но программирование на нём практически не масштабируется.

ThunderKittens занимает промежуточную позицию между Triton и CUDA C++. Он даёт явный контроль над иерархией памяти и tensor cores, но через абстракции, которые чувствуют себя естественно для ML-разработчика. Тайл st_bf<64,64> читается почти как torch.Tensor размером 64×64, но под капотом несёт всю информацию о swizzling, TMA-метаданных и fragment layout. Это позволяет исследователям, привыкшим к PyTorch, быстрее осваивать high-performance kernel development без необходимости сначала становиться CUDA-экспертами.

Реальный пример: attention prefill

Чтобы показать, как всё это работает вместе, автор статьи реализовал non-causal attention prefill на ThunderKittens с использованием шаблона lcf. Ядро определяет layout: глобальные тензоры Q, K, V и O; входной блок с K и V, которые протекают через pipeline; scratch-блок с Q, который остаётся в SMEM на протяжении всей задачи; и accumulator в регистрах для промежуточных результатов softmax.

Реализация занимает считанные десятки строк — вся сложность double-buffering, TMA-загрузок и warp-group синхронизации скрыта за pipeline engine. При этом производительность оказывается вполне конкурентоспособной: на H100 PCIe TK-ядро обгоняет FlashAttention-2 через PyTorch SDPA в среднем в 1.55 раза (44–59% выше throughput в зависимости от длины последовательности). На коротких последовательностях (N=768) оно даже опережает FlashAttention-3, хотя на длинных (N≥1536) FA3 восстанавливает лидерство с преимуществом 5–15%.

Важный нюанс: это не тот же kernel, что в оригинальной статье ThunderKittens. Автор использовал компактную реализацию на lcf-шаблоне, а не кастомное ядро из публикации. Даже в таком «облегчённом» виде результаты показывают, что абстракции TK не вносят катастрофических накладных расходов — а это главный аргумент в пользу DSL.

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

ThunderKittens заменяет Triton?

Нет, это разные уровни абстракции. Triton скрывает CUDA полностью и компилирует тайловые программы в оптимизированные ядра автоматически. ThunderKittens встроен в CUDA и даёт явный контроль над memory hierarchy, warp-координацией и tensor core instructions. Triton быстрее для прототипирования, TK — для случаев, где нужен точный контроль над железом.

Нужно ли знать CUDA, чтобы использовать ThunderKittens?

Да, базовое понимание GPU-иерархии (warps, shared memory, occupancy, tensor cores) необходимо. TK снижает порог входа по сравнению с raw CUDA, но не устраняет его полностью. Если вы не знаете, что такое warpgroup или почему bank conflicts важны, придётся разобраться.

На каком железе работает ThunderKittens?

Официально поддерживаются NVIDIA Hopper (H100) и Blackwell (B200). Абстракции спроектированы вокруг архитектурных особенностей этих поколений: wgmma-инструкций на Hopper и tcgen05 на Blackwell. Поддержка старых GPU (Ampere, Ada) ограничена или отсутствует.

Итог

ThunderKittens — это попытка найти баланс между двумя крайностями GPU-программирования: высокоуровневыми фреймворками, которые скрывают железо, и низкоуровневым CUDA, который требует экспертизы и времени. Тайловые абстракции, pipeline-шаблоны и прямой доступ к tensor cores позволяют писать kernels, которые конкурируют с hand-optimized реализациями, но при этом умещаются в десятки строк вместо сотен.

Для индустрии это означает, что цикл от алгоритмической идеи до production-ready kernel сокращается. А для scaling laws это означает, что второй множитель — FLOPS per Dollar — можно двигать не только за счёт нового железа, но и за счёт более эффективного использования существующего. В мире, где каждый процент throughput трансформируется в миллионы долларов на обучении, такие инструменты становятся стратегически важными.