Block-sparse ядра на GPU: когда нули стоят денег

Плотное матричное умножение на GPU не различает нули и значимые числа: оно честно перемножает всё подряд. Если 90% весов модели равны нулю, три четверти FLOPs уходит в никуда, а память тратится на хранение бесполезных элементов. Block-sparse ядра решают эту проблему грубо, но эффективно: они делят матрицу на блоки фиксированного размера (например, 16×16 или 32×32) и полностью пропускают блоки, состоящие из одних нулей. Ни загрузки, ни умножения, ни записи результата.

Ключевое слово здесь — block. Обычная неструктурированная разреженность (когда нули разбросаны по матрице как попало) на GPU почти не даёт ускорения: доступ к памяти становится нерегулярным, а видеокарта любит непрерывные, выровненные обращения. Блочная структура возвращает регулярность и позволяет ядру работать почти на пиковой пропускной способности — но только для тех блоков, которые остались.

Почему неструктурированная разреженность не спасает

GPU выполняет вычисления группами потоков (warp — 32 потока на архитектурах NVIDIA), которые исполняют одну инструкцию над разными данными. Когда нули лежат случайно, потоки внутри одной группы обрабатывают перемешанные значимые и нулевые элементы. Пропустить умножение для одного потока нельзя — остальные всё равно ждут. Плюс gather/scatter по разбросанным индексам убивает кэш и коалесцинг обращений к памяти.

В итоге модель, обрезанная до 80% нулей неструктурированным прунингом, на обычном GEMM* работает не быстрее плотной. Экономится только диск и, при специальном хранении, часть памяти. Скорость инференса — нет.

Разреженность, которую GPU не может прочитать регулярно, — это разреженность только на бумаге. Железо считает нули так же старательно, как единицы.

Что меняет блок

Если заставить нули собираться в блоки, warp целиком обрабатывает либо значимый блок, либо пропускает нулевой без ветвления внутри группы. Обращения к памяти становятся непрерывными. Отсюда и реальное ускорение, пропорциональное доле выкинутых блоков.

Форматы хранения

Классический CSR (Compressed Sparse Row) хранит отдельные ненулевые элементы и их индексы — он хорош для научных разреженных матриц, но плохо ложится на тензорные ядра. Для block-sparse используют блочные аналоги.

Особняком стоит формат 2:4: это не «крупные блоки», а мелкозернистая структурированная разреженность, которую с архитектуры Ampere умеют ускорять аппаратно. Ограничение жёсткое — ровно два ненулевых из каждой четвёрки. Теоретический потолок ускорения матричной части — двукратный, на практике меньше из-за накладных расходов и того, что не всё время модель проводит в GEMM.

Где это уже работает

Практических путей внедрить block-sparse сегодня несколько, и они сильно различаются по порогу входа.

• cuSPARSE / cuSPARSELt — библиотеки NVIDIA. cuSPARSELt заточена под structured sparse (2:4) и тензорные ядра, cuSPARSE поддерживает BSR и Blocked-ELL для более общих случаев.
• Triton — язык от OpenAI для написания GPU-ядер на Python-подобном синтаксисе. Позволяет писать собственные block-sparse матмулы, задавая размер блока и маску, без спуска в CUDA C.
• Готовые прунинг-пайплайны — например, инструменты, которые дообучают модель под шаблон 2:4, а затем экспортируют её для инференса на Sparse Tensor Cores.

Минимальный набросок ядра на Triton

Идея block-sparse матмула в Triton: итерируемся не по всем блокам K-измерения, а только по тем, что помечены в маске как ненулевые. Ниже — сильно упрощённый скелет, показывающий именно эту логику, а не готовую боевую реализацию.

import triton
import triton.language as tl

@triton.jit
def block_sparse_matmul(
    a_ptr, b_ptr, c_ptr,
    block_idx_ptr, n_blocks,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    # идём только по ненулевым K-блокам из индекса
    for i in range(n_blocks):
        k = tl.load(block_idx_ptr + i)
        offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
        offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
        offs_k = k * BLOCK_K + tl.arange(0, BLOCK_K)

        a = tl.load(a_ptr + offs_m[:, None] * stride_am
                          + offs_k[None, :] * stride_ak)
        b = tl.load(b_ptr + offs_k[:, None] * stride_bk
                          + offs_n[None, :] * stride_bn)
        acc += tl.dot(a, b)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    tl.store(c_ptr + offs_m[:, None] * BLOCK_N + offs_n[None, :], acc)

Суть в цикле по block_idx_ptr: он перечисляет индексы только тех K-блоков, что реально содержат данные. Плотное ядро прошло бы по всем блокам подряд. Точные сигнатуры и оптимизации (autotune, конвейеризация загрузок) сверяйте с текущей документацией Triton — API меняется между версиями.

Цена вопроса: когда овчинка стоит выделки

Block-sparse — не бесплатный тумблер «сделать быстро». За ускорение платят точностью и инженерным временем.

1. Дообучение под шаблон. Просто занулить веса по маске мало — модель просядет по качеству. Обычно нужен sparse-aware дообучающий проход, чтобы сеть подстроилась под ограничение вроде 2:4.
2. Порог полезной разреженности. При низкой доле нулей накладные расходы на индексы и маски съедают выигрыш. Крупные блоки дают больше ускорения, но грубее режут точность.
3. Не весь инференс — это GEMM. Ускорив матричные умножения вдвое, вы не ускорите модель вдвое: attention-софтмакс, нормализации, активации и передача данных остаются плотными.

Практический ориентир: если разреженность структурирована под 2:4 и модель дообучена, потеря качества часто удерживается в пределах долей процента при заметной экономии на матричной части. Но конкретные цифры зависят от архитектуры модели и задачи — их надо мерить на своём железе и своих данных, а не брать из чужого бенчмарка.

* GEMM (General Matrix Multiply) — базовая операция умножения двух матриц с накоплением, на которую опирается большая часть вычислений в нейросетях. Плотный GEMM обрабатывает все элементы, block-sparse — только ненулевые блоки.