Оптимизация GPU-ядер обычно требует глубокого понимания аппаратной архитектуры и низкоуровневых инструкций. Языковые модели пока справляются с этим нестабильно: они могут сгенерировать правдоподобный код, но далеко не всегда он оказывается корректным и быстрым. PTXBench — новый бенчмарк, который пытается внести в этот процесс измеримость и ясность.
Что такое PTXBench и зачем он нужен
PTXBench — это оценочный стенд для больших языковых моделей, которые пытаются использовать архитектурно-специфичный PTX (Parallel Thread Execution) для оптимизации GPU-ядер. PTX — промежуточное представление инструкций для GPU от NVIDIA, и именно на этом уровне можно тонко управлять регистрами, памятью и планировщиком. Обычные бенчмарки проверяют только итоговый код, но PTXBench смотрит глубже: действительно ли модель применяет целевые низкоуровневые инструкции, а не просто генерирует «похожий на C» код.
Авторы ориентировались на две классические нагрузки: умножение матриц (GEMM) и механизм внимания (attention), который критичен для трансформеров. Тесты выполнялись на актуальных ускорителях — H100 и B200. Такой выбор позволяет проверить не только базовые навыки модели, но и её способность адаптироваться к современным архитектурам с их специфическими особенностями.
Что измеряет бенчмарк
PTXBench оценивает модели по трём ключевым параметрам:
- Функциональная корректность — результат работы сгенерированного ядра должен совпадать с эталонным.
- Выполнение целевых инструкций — проверяется, встречаются ли в рантайме те самые PTX-инструкции, которые модель должна была использовать.
- Ускорение — насколько полученное ядро быстрее frontier-библиотек, то есть лучших готовых реализаций.
Такой подход разделяет три разные проблемы. Модель может написать корректный код, но при этом не использовать требуемые инструкции. Или использовать их, но получить производительность хуже стандартной библиотеки. PTXBench позволяет увидеть, на каком именно этапе происходит сбой.
Результаты: оптимизация пока даётся моделям неравномерно
Оценка показала, что способность LLM работать с PTX сильно зависит от задачи. На простых случаях модели демонстрируют неплохие результаты, но доля успешных решений резко падает на сложных backward-задачах для attention. Это ожидаемо: обратный проход требует более аккуратной работы с градиентами и памятью.
Ещё один важный вывод: выполнение целевых инструкций само по себе не гарантирует скорости. Модель может «попасть» в нужные PTX-инструкции, но выбрать неудачную стратегию распараллеливания или расстановку синхронизаций. В итоге ни одна из протестированных моделей не смогла стабильно конкурировать с frontier-библиотеками на всём наборе задач.
Как дообучали модель и что сработало
Чтобы понять, можно ли улучшить эти навыки, авторы провели дообучение модели Qwen3.6-27B с помощью supervised fine-tuning. Дополнительно они использовали приём, который можно назвать обучением с условием на исправление: модель учится не просто генерировать код, а исправлять его, получая обратную связь о неверных шагах.
Такой подход действительно улучшил несколько задач, но не обеспечил равномерного прогресса. Обобщение осталось слабым: модель могла хорошо справляться с одними типами ядер и проваливаться на других. Анализ показал, что решает не только размер обучающего датасета. Важнее его покрытие, баланс между разными сценариями и качество «учителя», который генерирует примеры для reasoning-цепочек. Если в данных перекос в сторону простых случаев, модель не научится сложности.
Главный итог
PTXBench даёт сообществу проверяемый тестовый стенд, на котором можно отслеживать прогресс LLM в оптимизации GPU-ядер. Он не просто выдаёт оценку «быстро/медленно», а показывает, на каком уровне рушатся планы модели: на синтаксисе, на применении специфичных инструкций или на выборе стратегии. Это важный шаг к тому, чтобы сделать автоматическую оптимизацию ядер предсказуемой и практичной.
