Написание GPU-ядер вручную требует терпения, усилий и, чего уж там, нервов. Я экспериментировал с современными LLM (Claude Opus 5, GPT 5.6 Sol), чтобы они буквально печатали для меня ядра. Расскажу, что я из этого вынес.
Если честно, написание GPU-ядер — это задача, которую до неприличия легко проверять.
Сначала ты решаешь, для каких операций будешь писать ядро. Затем пишешь первую версию и убеждаешься, что она компилируется без очевидных ошибок. Потом проверяешь корректность на медленной эталонной реализации. Если результат не сходится, пытаешься исправить проблему с корректностью. Когда это готово, ты бенчмаркаешь время выполнения ядра. Последующие версии строятся поверх этого, и ты продолжаешь оптимизировать, пока не достигнешь roofline-метрик или не удовлетворишься результатом.
Проверяемый цикл: разработка GPU-ядер
На изображении выше это показано как проверяемый цикл:
- Проверка компиляции — это быстрый локальный цикл (B <-> C), который выполняется ещё до того, как ты вообще задумываешься о корректности.
- Проверка корректности (D <-> E <-> F) — основной цикл проверки и вознаграждения. Именно эта часть делает написание ядер хорошей задачей для автоматизированной проверки, ведь у тебя есть эталон, с которым можно сверяться.
- Оптимизация (G -> H -> обратно к D) переиспользует тот же цикл проверки корректности для каждой новой версии: быстрое, но неправильное ядро ничего не стоит, поэтому корректность нужно проверять для каждой версии.
- Проверка roofline/удовлетворённости (I) — это внешний цикл, который решает, продолжать ли оптимизацию или остановиться.
Первая версия
Изображение выше — это всё ещё общий взгляд, а дьявол кроется в деталях. Нужно убедиться, что у нашего LLM-агента есть весь необходимый контекст, чтобы вообще начать писать хорошую первую версию.
Здесь в игру вступают CUDA-DSL. Triton, CuTeDSL и Tilelang — это те, с которыми очень легко начать работать на Python. Кривая обучения здесь не такая крутая, как в CUDA C++, однако абстракции в этих DSL могут ещё сильнее запутать нашего агента. Нужен способ передать агенту контекст этих абстракций.
Современные LLM уже умеют писать «хороший» Triton. Они отлично работают с абстракциями Triton даже без какого-либо контекста. Однако для других DSL, таких как CuTeDSL (который даёт гораздо больше контроля, чем Triton), я обнаружил, что наличие директории с контекстом, в которой агент может разобраться в абстракциях DSL, очень помогает.
Например, клонирование репозитория
NVIDIA cutlass
в директорию с контекстом — хороший способ дать агенту возможность искать абстракции, связанные с
Layout Algebra, Copy/GEMM atoms, иерархией памяти, примерами ядер
и так далее, при написании ядер на CuTeDSL.
По моему опыту, хорошая первая версия ядра компилируется без очевидных ошибок и проходит тест на корректность, о котором я расскажу ниже.
Тестирование, бенчмаркинг и профилирование
Когда агенту дан достаточный контекст, настоящее узкое место смещается на валидацию. Сама эталонная реализация и сверка с ней становятся всё важнее. Я называю этот этап тестированием корректности или просто тестированием. Скорость эталонной реализации не так важна, как её назначение. То, что ты собираешься измерять и проверять, — это то, что твой агент будет оптимизировать.
Обычно, когда вычисления не должны выполняться в пониженной точности, я измеряю максимальную абсолютную/относительную ошибку (MAE), среднеквадратичную ошибку (MSE/RMSE) и PSNR (пиковое отношение сигнала к шуму). Когда задействованы пониженные точности, я обычно измеряю PSNR и косинусное сходство (cossim).
То, как ты будешь запускать версии ядра на GPU, зависит от того, доступен ли GPU локально или через облако. В любом случае наш агент должен иметь возможность так или иначе получить доступ к своим результатам.
Я считаю описанную ниже методологию rung хорошим способом организовать N тестовых функций:
1def rung(name):2 def deco(fn):3 try:4 out = fn()5 results[name] = {"ok": True, **(out or {})}6 print(f"[{name}] ok " + " ".join(f"{k}={v}" for k, v in (out or {}).items()))7 except Exception as e:8 results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}9 print(f"[{name}] FAILED {type(e).__name__}: {e}")10 traceback.print_exc()11 return fn12 return deco
которую можно вызывать так:
1out = {}23@rung("pre-checks")4def _():5 run_pure_checks()6 run_dsl_checks()78@rung("run")9def _():10 out["o"] = custom_kernel(*inputs)11 torch.cuda.synchronize()12 return {"shape": tuple(out["o"].shape),13 "finite": bool(torch.isfinite(out["o"]).all())}
Для бенчмарк-ступени (rung) можно сделать несколько вещей:
- Замерять полное время выполнения ядра от начала до конца, чтобы оценить общее затраченное время
- Использовать внутриядерную трассировку для замера секций внутри ядра и выгружать их в вывод (с помощью кастомного трейсера или CUPTI)
- Сохранять сгенерированные IR, PTX, SASS и CUBIN в директорию dumps и позволять агенту их изучать
Последний пункт можно раскрыть немного подробнее. Иногда случается, что DSL в процессе lowering генерирует неоптимальный PTX (а затем и SASS), и ты находишь более удачные инструкции или формы, которые можно использовать вместо этого. Наш агент может читать текстовые файлы PTX/SASS и встраивать код более низкого уровня напрямую, не позволяя DSL обрабатывать неоптимальную часть. Опять же, здесь очень полезно передавать документацию по PTX в качестве контекста, чтобы агент мог по ней искать.
Последнее, что связывает всё воедино, — это профилирование. Если твой агент имеет доступ к CLI утилиты NCU (Nsight Compute Systems), ты можешь попросить его профилировать ядро и сгенерировать отчёт в рамках описанного выше проверяемого цикла обратной связи.
Заключительные мысли
«Значит, разработка GPU-ядер мертва?»
«Ну да, но вообще-то нет»
Да, потому что сложная часть — лейауты, индексация, абстракции и общая структура — во многом решается агентами с достаточным контекстом. Ты легко можешь сократить объём работы с 2–3 недель до 1–2 дней. Нет, потому что реальное узкое место сместилось с ядер на валидацию. Теперь нет одного-единственного способа действовать: чем лучше твой контекст и обвязка (harness), тем быстрее идёт процесс. Специализированные случаи выиграют ещё больше, и всё, что тебе понадобится, — это хорошая обвязка.
И наконец, вместо того чтобы относиться к агентам как к автономным, их всё ещё нужно воспринимать как по-настоящему умных ассистентов, которыми ты можешь управлять. Именно здесь пригодится твоё фундаментальное понимание GPU и ядер. Человеческая часть (ты) здесь всё ещё нужна.
Это горько-сладкое чувство, я знаю :)
