Hand-written PTX Tensor-Core GEMM kernels показывают, что выигрыш от низкоуровневого кода зависит не от “близости к железу”, а от того, что именно он убирает из пути выполнения. На NVIDIA L4 это особенно заметно для INT8 и INT4.
В этой работе сравниваются hand-written PTX Tensor-Core GEMM kernels с базовым WMMA-путём на одной GPU, NVIDIA L4 на Ada SM89. Важен не сам факт ускорения, а причина, по которой оно появляется или исчезает. Авторы проверяют FP16, INT8 и INT4 на размерах от N = 512 до N = 8192 и смотрят не только на wall time, но и на поведение памяти, кэшей и счетчиков исполнения.
Проблема начинается там, где WMMA перестает быть удобной абстракцией и становится ограничением. Высокоуровневый API фиксирует shape фрагментов, layout загрузки и раскладку аккумуляторов. Это делает код проще, но оставляет меньше пространства для оптимизации. PTX даёт прямой контроль над cp.async, ldmatrix и mma.sync, но за это приходится платить сложностью, ручным управлением тилеями и риском сделать kernel хуже, а не лучше.
Авторы отвечают на практический вопрос: когда замена WMMA на PTX действительно дает end-to-end выигрыш? Их подход прагматичен. Они держат одинаковую block- и warp-tiling схему, сравнивают каждый PTX kernel только с WMMA-базой той же точности и меняют по одному параметру за раз. Это хороший инженерный ход, потому что он позволяет связать результат с конкретной причиной, а не с набором случайно улучшенных деталей.
Для FP16 вывод оказался простым: ручной PTX не побеждает compiler-optimized WMMA baseline. На малых и средних размерах kernel находится в compute-bound режиме, но даже там выигрыш от более низкоуровневых инструкций съедается overhead на packing operands. На больших размерах система упирается в память. Когда рабочий набор выходит за пределы L2, различия между реализациями почти исчезают. Иными словами, если bottleneck уже в DRAM, оптимизация instruction stream не меняет итоговую картину.
INT8 показывает более интересный компромисс. Лучший kernel, int8_ptx_mma_k32, стабильно быстрее WMMA на всех размерах. Причина не в магии Tensor Cores, а в более компактном instruction count и почти идеальном global-load coalescing. Авторы отдельно показывают, что occupancy здесь не объясняет throughput: самый высокий occupancy у WMMA-ядра, но оно не самое быстрое. Это важное наблюдение для SRE и performance-инженеров. Высокая occupancy может быть хорошим сигналом, но не всегда является целевой метрикой.
INT4 дает самый большой эффект. Здесь выигрыш связан не просто с более удачным расписанием, а с тем, что WMMA-путь для s4 фактически software-expanded. То есть он несет лишние инструкции и lane-dependent divergence. Native PTX kernel использует mma.sync.m16n8k64.s4 напрямую, и это снимает слой эмуляции. В результате появляются speedups в диапазоне 2.9×–4.3× относительно INT4 WMMA, а относительно FP16 WMMA лучшие quantized kernels доходят до 98.7× на N = 8192. Но и здесь картина не линейная: для малых N лучше трехстадийный pipeline, а для больших N выигрывает k64-вариант, который лучше сохраняет L1 locality.
Отдельный интерес представляет абляция вокруг INT4 k64 family. Авторы проверяют split loader, cache-eviction policy и layout B-operand. Итог аккуратный: выбранная конфигурация уже находится в local optimum. Изменения policy или layout не дают улучшения, а transposed B резко ухудшает coalescing и срезает usable bandwidth. Это полезный результат именно потому, что он убирает иллюзию “ещё одного простого тюнинга”. Иногда оптимальная точка уже найдена, и следующий шаг только перераспределяет pressure между L1, L2 и issue slots.
Главный вывод статьи хорошо ложится в инженерную практику. PTX оправдан, когда он убирает накладные расходы, которые WMMA не может спрятать: эмуляцию, лишние инструкции, плохой coalescing. Для FP16 это не срабатывает. Для INT8 это даёт умеренный, но стабильный выигрыш. Для INT4 это уже структурное преимущество native MMA path. При этом occupancy остаётся диагностическим сигналом, а не целью. Для inference на quantized LLM на одной L4-class GPU авторы рекомендуют сначала смотреть на memory-system behavior, а потом уже на Tensor Core utilization.
Источник информации
arXiv — крупнейший открытый репозиторий препринтов (с 1991 года, под эгидой Корнелла), где исследователи оперативно размещают рабочие версии статей; материалы общедоступны, но не проходят полное рецензирование, поэтому результаты следует считать предварительными и, по возможности, сверять с обновленными версиями или рецензируемыми журналами. arxiv.org