Система обработки нейросетевых моделей становится все более сложной и ресурсоемкой, что требует непрерывного улучшения программного обеспечения и аппаратной оптимизации. Одним из ключевых компонентов, определяющих эффективность вычислений в современной глубоковычислительной технике, являются вычислительные ядра, или kernel'ы — базовые алгоритмы, которые реализуют операции над тензорами. Особенно это важно для задач с гигантскими объемами данных и моделей с миллиардами параметров. В этой связи создание кастомных ядер, специально адаптированных под архитектуру конкретных GPU, становится важным направлением. Именно на этом сосредоточено внимание при разработке ядер для AMD MI300 — мощного графического процессора нового поколения, способного раскрыть потенциал больших моделей искусственного интеллекта.
AMD MI300 представляет собой весьма сложное устройство, состоящее из сотен вычислительных блоков и множества уровней памяти, что требует глубокого понимания его архитектуры для эффективной оптимизации вычислительных задач. Архитектура MI300 включает в себя 8 ускорительных чиплетов — XCD (Accelerator Complex Dies), объединяющих в себе 38 вычислительных блоков (Compute Units, CU) каждый. Всего на GPU насчитывается около 304 CU, способных одновременно исполнять огромное количество потоков. Такой параллелизм позволяет запускать сотни тысяч потоков одновременно, что дает интеграторам возможность решать комплексные задачи с оптимальной пропускной способностью. Однако высокая параллельность требует ограничений и координации, поскольку потоки на уровне варпа, состоящего из 64 потоков, должны выполнять одну и ту же инструкцию синхронно.
Кроме того, MI300 оснащен несколькими уровнями кэш-памяти, включая L1 на каждом CU, L2, общую для XCD, и огромный Infinity Cache объемом 256 МБ, позволяющий снижать задержки при доступе к данным. Кроме того, графический процессор располагает огромным объемом видеопамяти 192 ГБ, что особенно ценно для обработки крупных моделей. Одной из главных особенностей работы на MI300 является необходимость минимизировать обращения к глобальной видеопамяти из-за высокой латентности. Такое требование заставляет разработчиков писать ядра, которые максимально эффективно используют кэш и совместное использование памяти на уровне блока потоков (shared memory). Создание кастомных ядер становится необходимостью, поскольку стандартные библиотеки оптимизированы, как правило, под более универсальные задачи, а не под узкоспециализированные сценарии, например, специфику инференса крупномасштабных моделей на FP8 с определенной геометрией тензоров.
Одним из главных примеров значительного увеличения производительности является оптимизация ядра RMS нормализации с гибридной операцией усреднения, нормализации и конвертации в FP8. Традиционно эти операции выполняются раздельно, вызывая значительные затраты времени. Интеграция их в единый конвейер с грамотным распределением потоков и использованием синхронизации на уровне блока потоков позволила добиться до 10-кратного прироста производительности по сравнению с базовыми реализациями в PyTorch. Основу успеха составляет грамотное распределение вычислительных ресурсов и минимизация операций записи и чтения из глобальной памяти. Особое внимание уделяется загрузке данных шириной 128 бит за одну инструкцию, обеспечивающей скоординированный, непрерывный доступ к памяти всеми потоками в варпе.
Такой подход снижает количество необходимых циклов загрузки и, следовательно, время ожидания. Важной оптимизацией стало использование общей памяти (shared memory) для промежуточного хранения данных, что сводит к минимуму необходимость повторного обращения к более медленной видеопамяти. Аналогичный подход был применен к ядру, реализующему функцию активации SwiGLU, которая используется в блоке MLP. Оптимизация этого ядра включала использование векторизованных инструкций для FP16, позволяющих сразу обрабатывать несколько элементов, что существенно увеличивает пропускную способность и снижает количество выполняемых инструкций. Кроме того, замена операции экспоненты (exp) на более быструю exp2 с предварительным масштабированием входных данных позволила снизить задержки при вычислении функции сигмоидной активации без заметных потерь точности.
Существенную часть времени при выполнении инференса занимают операции свертки и матричного умножения общего вида (GEMM). Несмотря на высокий уровень оптимизации стандартных библиотек hipBLAS и rocBLAS под универсальные задачи, ситуация меняется при работе с «тонкими» GEMM — с небольшим числом строк. Такие задачи вызывают недостаточную нагрузку на вычислительные блоки GPU, что снижает общую эффективность. Для решения этой проблемы была реализована техника split-K, позволяющая разбивать GEMM по общей размерности K на несколько параллельных подзадач. Это увеличивает число одновременно загруженных в вычисление блоков и, таким образом, повышает загрузку GPU.
Для максимизации производительности также был применён метод специализированного распределения варпов: части варпов выделяются для загрузки данных (продюсеры), а другие — для вычислений (консюмеры). Такой асинхронный подход распределения работы позволяет лучше скрывать задержки обращения к памяти и повысить арифметическую интенсивность задачи. Важным улучшением стала поддержка 4:2 разрежённости ввода в специальных MFMA (Matrix Fused Multiply-Add) инструкциях. Это позволило эффективно использовать аппаратную поддержку разрежённых операций, характерную для MI300, а также минимизировать добавочные расходы памяти за счет устранения лишних операций с нулями. Подобный подход обеспечивает заметный выигрыш в производительности на низких размерах входных данных, что актуально для параллельной обработки запросов моделью с небольшими батчами.
При интеграции оптимизированных кастомных ядер в систему управления инференсом VLLM и окружение с 8 GPU MI300X удалось достичь значительного улучшения времени отклика — критически важного параметра при масштабировании сервисов с миллионами запросов. Анализ показал, что ядра для операций RMS нормализации и SwiGLU активации снизили их долю в общей задержке на 15%, а новый skinny GEMM kernel обеспечил ускорение при специфических проекциях модели. Совместное использование этих подходов привело к общей суммарной оптимизации латентности почти по всем ключевым этапам вычислений. Полученные в результате ускорения метрики значительно увеличивают эффективность работы нейросетевых моделей на AMD MI300, делая эту платформу конкурентоспособной с решениями Nvidia. Благодаря открытому доступу к исходкам созданных ядер в репозитории hf-rocm-kernels сообществу разработчиков предоставляется возможность воспроизвести результаты, адаптировать и расширить реализованные оптимизации под собственные задачи.
