Агент Codex прискорив ядро QR у 232 рази в конкурсі з авто-досліджень

У конкурсі з авто-досліджень, який організували GPU Mode та Core Automation, розробник під псевдонімом sankalp посів 12-те місце серед 183 учасників. Його рішення працювало у 232 рази швидше за базове ядро QR-розкладу. У дописі в блозі він розповідає, як використав агента Codex від OpenAI та щільний цикл подань — підхід, який сам називає «loop engineering».

Задача: пакетне QR-розкладання Хаусхолдера

Учасники мали реалізувати пакетне квадратне компактне QR-розкладання Хаусхолдера. На вході — пакет матриць FP32 CUDA розміру batch × n × n, на виході — той самий компактний формат, який повертає torch.geqrf: верхній трикутник H зберігає R, нижня частина — вектори Хаусхолдера, а окремий вектор tau містить коефіцієнти відбиття. Перевіряльник відновлював Q через torch.linalg.householder_product(H, tau) і перевіряв умови A ≈ QR, QᵀQ ≈ I та QᵀA ≈ R.

Правильні подання ранжувалися за геометричним середнім часом виконання для різних розмірів і випадків обумовленості. Ключовими були розміри 512×512, а також 1024, 2048 і 4096. Внутрішньо дозволялися обчислення у FP16, FP8 чи NVFP4, але результат мав проходити перевірки рівня FP32.

Codex у циклі

GPU Mode надав CLI-інструмент popcorn, який дозволяв агентам тестувати, вимірювати та надсилати рішення на лідерборд. За 14 днів автор зробив понад 1500 подань. У робочому просторі Codex були опис задачі, AGENTS.md з інструкціями та журнал спроб. Для кількісних цілей він використовував команду /goal, для перевірки без зупинки циклу — /btw, а вузькі місця шукав за допомогою профілювання Modal і NCU.

Щоб не застрягти в локальних максимумах, автор тримав «пучок» із трьох-п'яти ідей-кандидатів, залучав сильнішу модель-радника та окремих субагентів для експериментів.

Від 419 мс до 1,8 мс

Базовий шлях torch.geqrf займав приблизно 419 000 мікросекунд. Після впровадження блокового алгоритму Хаусхолдера для n = 512 автор за день досяг 5000 мікросекунд, а фінальний результат становив 1805 мікросекунд. Блоковий підхід стискає відбиття панелі в одне оновлення WY рангу b, перетворюючи роботу над хвостовою частиною на три множення матриць — саме ту форму, яку найкраще обробляють тензорні ядра.

Висновки

Автор підсумовує, що галузева експертиза прискорює і побудову середовища, і керування моделлю: що краще розумієш предмет, то точніші запитання ставиш агенту. Це була його перша серйозна спроба авто-досліджень, і другий конкурс серії — з власних розкладів — уже триває.