Vibecoding GPU Kernels

Написання 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, ієрархією пам'яті, прикладами ядер,

тощо, під час написання ядер на CuTeDSL.

За моїм досвідом, хороша перша версія ядра компілюється без очевидних помилок і проходить тест на правильність, про який я розповім нижче.

Тестування, бенчмаркінг і профілювання

Коли в агента достатньо контексту, справжнім вузьким місцем стає валідація. Сама еталонна реалізація, а також перевірка проти неї стають дедалі важливішими. Я називаю цю фазу тестуванням правильності або просто тестуванням. Швидкість еталонної реалізації менш важлива, ніж її призначення. Те, що ви плануєте вимірювати та перевіряти, — саме те ваш агент і буде оптимізувати.

Зазвичай, коли обчислення не мають виконуватися зі зниженою точністю, я вимірюю максимальну абсолютну/відносну похибку (MAE), середньоквадратичну похибку (MSE/RMSE) та PSNR (пікове відношення сигналу до шуму). Коли задіяна знижена точність, я зазвичай вимірюю PSNR та косинусну подібність (cossim).

Те, як версії ядра реально запускаються на GPU, залежить від того, чи GPU доступний локально, чи через хмару. Попри це, наш агент повинен мати можливість так чи інакше отримувати доступ до своїх результатів.

Я вважаю наведену нижче методологію rung зручним способом мати будь-яку кількість тестових функцій:

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 fn
12    return deco

її можна викликати так:

1out = {}
2
3@rung("pre-checks")
4def _():
5    run_pure_checks()
6    run_dsl_checks()
7
8@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())}

Що стосується бенчмарк-рангу, можна зробити кілька речей:

• Бенчмаркати наскрізний час виконання ядра (end-to-end), щоб отримати загальний витрачений час
• Використовувати внутрішньоядерне трасування, щоб бенчмаркати окремі секції ядра та виводити їх у результатах (за допомогою власного трейсера або 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, то швидший процес. Спеціалізовані випадки виграють ще більше, і все, що вам знадобиться, — це хороший harness.

Нарешті, тут усе ще потрібно ставитися до агентів не як до автономних систем, а як до справді розумних помічників, якими ви можете керувати. Саме тут вам стане в пригоді фундаментальне розуміння GPU та ядер. Людська складова (ви) тут усе ще потрібна.

Це гірко-солодке відчуття, я знаю :)