Сначала 12-го февраля Nvidia выпустила блогпост. В нём рассказывается, как DeepSeek R1 может использоваться для оптимизации кернелов — это такие маленькие программы, выполняющие несколько простых функций на видеокарте (GPU). Кернелы используются для запуска LLM, многие (все?) отдельные операции или блоки операций прописаны как раз в них.
Создание оптимизированного кернела — то есть такого, который утилизирует GPU эффективно, на полную — требует специальные навыки и занимает много времени даже у опытных инженеров. Нужно думать о том, как обращаться с памятью GPU, в какие моменты где какие переменные запоминать, в каком порядке операции выполнять, в какой части памяти данные хранить — от всего этого зависит скорость. Иногда даже имеет смысл не сохранять промежуточный результат, а пересчитывать его после при необходимости (потому что память ценнее). Ну и вишенка на торте — для разных поколений GPU кернелы могут отличаться — потому что на ново архитектуре эффективнее посчитать что-то можно по-другому (пример FlashAttention 3). В общем, много сложных неочевидных решений.
Идея: давайте для каждого кернела напишем тесты, которые позволяет проверить, что он работает в точности так, как мы ожидаем. Если входные аргументы одни, то результат вот такой, а если вот такие — то другой. Теперь когда такие тесты есть, мы можем брать любую программу, прогонять через тесты и проверять, совпадает ли она с нашим ожиданием. А затем берём LLM, даём ей описание кернела, его код и говорим «сделай лучше». Моделька думает, выплёвывает код кернела, мы его тестируем на реальной GPU и замеряем время. Если тесты не прошли — отдаём модели на переработку, если прошли, то добавляем в промпт время (ну и тут понятно, стало лучше/хуже).
Повторяя эту операцию в цикле мы можем рассчитывать на некоторые улучшения (если LLM достаточно умна или кернелы были совсем неоптимизированными).
Nvidia хвастается, что они смогли ускорить несколько достаточно популярных кернелов (используемых в PyTorch FlexAttention, если вам интересно) на 10-110%. Это не означает, что все модели и всегда будут быстрее, так как отдельные кернелы используются в отдельных моделях. Плюс, кернел это очень малая часть вычислений — ускорение одного даже на 400% может привести к суммарному ускорению всей модели на, условно, 5%.
На картинке 1 — ускорения для разных кернелов относительно того, что сейчас есть в PyTorch (то есть по идее со следующим обновлением все должны получить ускоренные версии кернелов), а на картинке 2 — доля кернелов из некоторого бенчмарка, для которых модель смогла написать код, проходящий все тесты (без условия «он быстрее того что было»). Видно, что чем больше времени Nvidia давали подумать рассуждающей модели R1, тем больше эта доля.