Как CUDA, ROCm, Triton и TensorRT формируют производительность GPU в AI: компилеры и настройка

Пропускная способность в глубоких нейросетях определяется тем, как стек компиляторов отображает тензорные программы на исполнение на GPU: расписание потоков и блоков, перемещение данных и выбор инструкций для матричных конвейеров. Ниже сравнение четырех доминирующих стеков с точки зрения компиляции и практической настройки, и какие оптимизации действительно имеют значение.

Какие факторы влияют на производительность современных GPU

У разных вендоров повторяются одни и те же рычаги:

• Планирование и фьюзинг операторов: уменьшение числа запусков кернелов и обращений к HBM позволяет цепочкам producer->consumer переиспользовать регистры и shared memory. Runtime-фьюзинг в библиотеках особенно важен для блоков attention и сверточных цепочек.
• Тайлинг и раскладка данных: формы тайлов должны соответствовать нативным фрагментам Tensor Core или аналогичным матричным конвейерам; нужно избегать конфликтов в shared memory и partition camping. Руководства по warp-level GEMM дают канонические примеры.
• Точность и квантизация: FP16/BF16/FP8 для тренировки и inference; INT8/INT4 через калибровку или QAT для продакшен-инференса. Инструменты билда автоматизируют калибровку и выбор кернелов под эти точности.
• Захват графа и специализация во время выполнения: захват графов амортизирует оверхед запусков и позволяет динамический фьюзинг общих подграфов, например attention.
• Автотюнинг: поиск размеров тайлов, факторов распаковывания и глубин пайплайнинга под конкретную архитектуру и SKU часто решает исход производительности.

CUDA: nvcc/ptxas, cuDNN, CUTLASS и CUDA Graphs

Путь компиляции. CUDA-код проходит через nvcc в PTX, затем ptxas понижает PTX до SASS, специфичного машинного кода. Контроль оптимизаций требует передачи флагов для хостовой и девайсной фаз; для кернелов ключевые опции идут через -Xptxas. Учтите, что хостовый -O3 сам по себе не влияет на девайсную фазу.

Генерация кернелов и библиотеки. CUTLASS предоставляет параметризуемые шаблоны для GEMM и conv с warp-level тайлингом, Tensor Core MMA конвейерами и итераторами shared memory, спроектированными для доступа без конфликтов. cuDNN 9 добавил runtime-фьюзинг и интеграцию с CUDA Graphs, что существенно снижает оверхед диспетчеризации в задачах Transformer.

Последствия для производительности. Замена неперемешанных последовательностей PyTorch-операций на cuDNN attention fusion обычно сокращает число запусков кернелов и глобальный трафик памяти; в паре с CUDA Graphs это снимает узкие места на CPU при коротких последовательностях инференса. На Hopper и Blackwell согласование форм тайлов с нативными размерами WGMMA/WMMA критично: неправильно подобранные тайлы приводят к потере пропускной способности тензор-ядр. CUDA подходит, когда нужен максимальный контроль над выбором инструкций, occupancy и choregraphy shared memory или когда вы расширяете кернелы за пределы библиотек на NVIDIA.

ROCm: HIP/Clang, rocBLAS/MIOpen и серия 6.x

Путь компиляции. ROCm использует Clang/LLVM для компиляции HIP в ISA AMD. Версии 6.x сконцентрированы на производительности и покрытии фреймворков; релиз-ноты детализируют оптимизации и поддержку HW/OS.

Библиотеки и кернелы. rocBLAS и MIOpen реализуют примитивы GEMM и conv с учётом архитектуры и алгоритмического выбора, по духу схожие с cuBLAS и cuDNN. ROCm улучшает поддержку Triton на AMD, позволяя писать кернелы на Python и понижать их через LLVM в AMD-бекенд.

Последствия для производительности. На AMD важно согласование ширин банков LDS и векторизованных глобальных загрузок с формой тайлов матриц так же, как и выравнивание shared memory на NVIDIA. Компиляторный фьюзинг в фреймворках и автотюнеры в rocBLAS/MIOpen часто закрывают большую часть разрыва с ручными кернелами, при условии подходящей архитектуры и драйвера. ROCm стоит выбирать при необходимости нативной оптимизации под AMD и портировании из CUDA через HIP.

Triton: DSL на Python и компилятор для кастомных кернелов

Путь компиляции. Triton это DSL в Python, которое понижает код через LLVM, управляет векторизацией, коалесценцией памяти и распределением регистров, при этом даёт явный контроль над блок-тайлами и program-id. Triton автоматизирует трудоёмкие части CUDA-оптимизаций, оставляя автору выбор тайлинга блоков.

Оптимизации и автотюнинг. Triton поддерживает автотюнинг по размерам тайлов, числу warps и стадиям пайплайна; использует статическое маскирование, staging в shared memory и софтверный пайплайнинг для перекрытия загрузок и вычислений.

Последствия для производительности. Triton хорош для фьюзированных, shape-специализированных кернелов вне покрытия стандартных библиотек: кастомные attention, цепочки normalization-activation-matmul и т.п. На современных NVIDIA архитектурах совместная оптимизация снизила разрыв с CUTLASS-стилем кернелов. Triton выбирают, когда нужен практически CUDA-перформанс для кастомных операций без написания SASS/WMMA и важна быстрая итерация на Python.

TensorRT и TensorRT-LLM: оптимизация графов на этапе билда для inference

Путь компиляции. TensorRT принимает ONNX или графы из фреймворков и генерирует hardware-specific engine. На этапе билда происходят фьюзинг слоёв и тензоров, калибровка точности для INT8 или FP8/FP16 и подбор тактик кернелов. TensorRT-LLM добавляет оптимизации, ориентированные на большие языковые модели.

Оптимизации. Билд включает константный фолдинг, канонизацию concat-slice, фьюзинг conv-bias-activation и фьюзинг attention. Поддерживаются пост-тренировочная калибровка, smooth-quant и QAT-прописы. На рантайме доступны paged KV cache, in-flight batching и планирование для multi-stream и multi-GPU развертываний.

Последствия для производительности. Основные выигрыши даёт сквозной INT8 или FP8 там, где поддерживается, удаление фреймворкового оверхеда через один engine и агрессивный фьюзинг attention. TensorRT генерирует планы под конкретную архитектуру, избегая универсальных кернелов в рантайме.

Практические советы: выбор и настройка стека

• Тренировка и экспериментальные кернелы: CUDA + CUTLASS на NVIDIA или ROCm + rocBLAS/MIOpen на AMD; Triton для кастомных фьюзированных операций.
• Продакшен-инференс на NVIDIA: TensorRT или TensorRT-LLM для глобальной оптимизации графа и квантизации.
• Используйте нативные инструкции архитектуры: маппьте тайлы под WGMMA/WMMA на Hopper/Blackwell, согласовывайте LDS и векторные ширины на AMD.
• Сначала фьюзинг, затем квантизация: фьюзинг снижает трафик памяти, квантизация увеличивает плотность вычислений; сочетание часто дает мультипликативный выигрыш.
• Применяйте graph execution для коротких последовательностей: CUDA Graphs + cuDNN attention амортизируют оверхед запусков при autoregressive inference.
• Относитесь к флагам компилятора как к первоклассным параметрам: для CUDA не забывайте device-side флаги, например -Xptxas -O3,-v и -Xptxas -O0 для отладки; хостовый -O3 недостаточен для оптимизации девайсной фазы.

Для дальнейшего чтения и ссылок изучайте материалы NVIDIA, AMD, Triton и репозитории CUTLASS, rocBLAS, MIOpen и TensorRT.