Умножение матриц: от определения до кэш-линий

Каждая модель, которую я выводил в продакшен — каждый слой трансформера, каждая value-голова, каждый речевой энкодер — тратит большую часть своих FLOP внутри одной операции. Для $A \in \R^{m \times k}$ и $B \in \R^{k \times n}$ произведение $C = AB$ определяется поэлементно:

c_{ij} = \sum_{p=1}^{k} a_{ip}\, b_{pj}, \qquad 1 \le i \le m,\ \; 1 \le j \le n.

Это $mnk$ операций умножения-сложения — $2mnk$ операций с плавающей запятой — и определение дословно переводится в первый matmul, который пишет каждый. Определение — это три цикла. А производительность — это история про память, и этот пост спускается по ней до самой кэш-линии.

Три способа потратить 268 миллионов FLOP

При $m = n = k = 512$ произведение стоит $2 \cdot 512^3 \approx 2.7 \times 10^8$ FLOP. Вот определение, дословно — именно эту функцию замеряет скрипт бенчмарка:

def matmul_naive(A, B, n):
    C = [[0.0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            acc = 0.0
            for p in range(n):
                acc += A[i][p] * B[p][j]
            C[i][j] = acc
    return C

На моём десктопе это занимает 9,3 с на списках Python. Очевидная «оптимизация» — те же три цикла по массивам NumPy со скалярной индексацией (A[i, p] * B[p, j]) — занимает 41,9 с, в 4,5 раза дольше простых списков. А A @ B, уходящий в OpenBLAS, — 5,1 мс в один поток. Те же матрицы, те же 268 миллионов FLOP, ответы совпадают с точностью до $2.3 \times 10^{-13}$ :

Сравнение в лог-шкале одного умножения 512³: скалярная индексация NumPy — 41,9 с, списки Python — 9,3 с, OpenBLAS — 5,1 мс — рис. 1 — одно умножение, три реализации; обратите внимание на логарифмическую шкалу

Две вещи на этом графике требуют объяснения. Скандал — NumPy проигрывает обычным спискам в 4,5 раза — объясняется просто: каждый A[i, p] пересекает границу C-API и упаковывает свежий Python-объект float, около 120 нс накладных расходов на каждое обращение к элементу, 268 миллионов раз. Контракт NumPy — операции над целыми массивами; индексируйте его как список — и вы платите за всю машинерию, ни разу её не задействовав.

Интересное число — второе: 1800× между честными циклами и A @ B, притом что OpenBLAS выполняет ровно те $2mnk$ FLOP, которых требует определение, в той же двойной точности. Интерпретатор отвечает примерно за два порядка из этого разрыва — скомпилированная -O3-версия того же тройного цикла обычно выдаёт один-два GFLOP/s на ядре этого класса. Остальное — множитель в тридцать с лишним, переживающий компиляцию, — это память.

Окружение бенчмарка

Intel i7-6700K (Skylake, 4 ядра / 8 потоков, 4,0 ГГц база / 4,2 ГГц турбо на одном ядре), кэши 32 КБ L1d + 256 КБ L2 на ядро, 8 МБ общего L3, двухканальная DDR4. Python 3.11.7, NumPy 2.4.2 с scipy-openblas, всюду float64. BLAS ограничен одним потоком через OPENBLAS_NUM_THREADS=1, кроме многопоточных замеров в конце. Два медленных цикла запускаются по одному разу; все остальные замеры — лучший из 5 с фиксированным сидом. Результаты сверены с A @ B: максимальное абсолютное отклонение ≈ 2,3 × 10⁻¹³.

64-байтовая правда

CPU никогда не читает из памяти один float64. Он читает кэш-линию — 64 байта, восемь double — и держит её в иерархии кэшей: на этой машине 32 КБ L1d на ядро, 256 КБ L2, 8 МБ L3 на всех. Загрузка из L1 стоит ~4 такта; поход в DRAM — пару сотен. Всё в быстром численном коде следует из одного правила: вытащил линию — используй то, что на ней, и переиспользуй до того, как её вытеснят.

NumPy хранит матрицы построчно (C-порядок): строка $i$ непрерывна, а элемент $b_{pj}$ лежит в 8 байтах после $b_{p,j-1}$ , но на целую строку дальше, чем $b_{p-1,j}$ . Теперь посмотрите на внутренний цикл определения: он идёт по A[i][p] вдоль строки — последовательно, восемь полезных double на каждую линию, паттерн, который аппаратный префетчер распознаёт и начинает подкачивать данные с опережением, — и по B[p][j] вниз по столбцу, прыгая на 4 КБ за шаг. Каждый шаг этого спуска открывает новую кэш-линию и новую страницу памяти; один проход по столбцу трогает 512 разных линий (32 КБ трафика ради 4 КБ полезных данных), и префетчеру не за что зацепиться.

Эффект легко изолировать вообще без matmul. Возьмём одну матрицу 8192×8192 — 537 МБ, сильно больше любого кэша — и просуммируем её дважды: раз по строкам, раз по столбцам:

M = rng.standard_normal((8192, 8192))

s = 0.0
for i in range(8192): s += M[i, :].sum()   # по строкам:    44 мс

s = 0.0
for j in range(8192): s += M[:, j].sum()   # по столбцам: 597 мс

Те же байты, те же сложения — разница в 13,6 раза. Проход по строкам прокачивает целые кэш-линии вслед за префетчером. Проход по столбцам использует один double с каждой открытой линии за проход, а его шаг в 64 КБ трогает 8192 разных страницы памяти на столбец — сильно больше, чем вмещает TLB, так что обращения платят за обход таблиц страниц поверх кэш-промахов. (Соседние столбцы позже добирают остальные семь double каждой линии из L3 — потому штраф 13,6×, а не хуже.) Проход по строкам перекачивает 537 МБ за 44 мс — около 12,2 ГБ/с, практическая пропускная способность чтения этого ядра. Запомните это число.

Арифметическая интенсивность, или почему наивный цикл не может быть быстрым

Ядро ограничено двумя потолками: как быстро оно считает и как быстро его кормят. Это ядро на 4,2 ГГц с двумя 256-битными FMA-портами упирается в

4.2\ \text{ГГц} \times 2\ \tfrac{\text{FMA}}{\text{такт}} \times 4\ \tfrac{\text{double}}{\text{FMA}} \times 2\ \tfrac{\text{FLOP}}{\text{double}} \;\approx\; 67\ \tfrac{\text{GFLOP}}{\text{s}}

в двойной точности — но выдерживает лишь ~12,2 ГБ/с чтения. Какой потолок действует, решает арифметическая интенсивность: FLOP на каждый перемещённый байт.

Наивный цикл — на размерах, где операнды переросли кэши, — прогоняет оба операнда через ядро по разу на каждое использование: два 8-байтовых чтения на умножить-сложить, интенсивность $q = \tfrac{2\ \text{FLOP}}{16\ \text{Б}} = 0.125$ . При 12,2 ГБ/с это ограничивает даже идеально векторизованный цикл с таким паттерном доступа потолком примерно в $0.125 \times 12.2 \approx 1.5$ GFLOP/s — в сорок с лишним раз ниже вычислительного пика, ещё до единого такта интерпретатора. Этот потолок компиляция не поднимает: цикл сам по себе упёрт в память по построению.

Чтобы стать быстрее, не нужно меньше FLOP. Нужно больше FLOP на байт.

Блочное разбиение: идея √M

Решению не один десяток лет, и на нём до сих пор держится каждый BLAS и каждое GPU-ядро матричного умножения: разрежь задачу на тайлы так, чтобы небольшое рабочее множество жило в кэше и переиспользовалось до вытеснения.

Режем матрицы на тайлы $b \times b$ и накапливаем каждый тайл $C$ как сумму маленьких тайловых произведений:

for i0 in range(0, n, b):
    for j0 in range(0, n, b):
        # этот тайл C переиспользуется весь p-цикл
        for p0 in range(0, n, b):
            C[i0:i0+b, j0:j0+b] += A[i0:i0+b, p0:p0+b] @ B[p0:p0+b, j0:j0+b]

Один шаг внутреннего цикла трогает три тайла $b \times b$ — $3b^2 \times 8$ байт — и выполняет над ними $2b^3$ FLOP:

q(b) \;=\; \frac{2b^3}{24\, b^2}\ \tfrac{\text{FLOP}}{\text{Б}} \;=\; \frac{b}{12}\ \tfrac{\text{FLOP}}{\text{Б}},

— интенсивность, которая растёт с размером тайла. Чтобы поднять это ядро из memory-bound в compute-bound, нужно $q \gtrsim 67 / 12.2 \approx 5.5$ , то есть $b$ около семидесяти — а три тайла 70×70 в float64 занимают 118 КБ и спокойно помещаются в 256 КБ L2. В этом весь трюк, и именно поэтому ёмкость кэша $M$ по всей этой литературе ходит под квадратным корнем: тайлы со стороной $b \sim \sqrt{M/3}$ срезают суммарный трафик с наивных $\mathcal{O}(n^3)$ слов до

\mathcal{O}\!\left(\frac{n^3}{\sqrt{M}}\right),

что, как доказали Хонг и Кунг в 1981 году, асимптотически оптимально для любого порядка выполнения классического алгоритма. Больше кэш — меньше трафик, причём выигрыш растёт как квадратный корень из ёмкости.

Что добавляет настоящий BLAS

OpenBLAS — это та же идея, исполненная с одержимостью: размеры тайлов подобраны под каждый уровень кэша, packing операндов (каждый тайл копируется в непрерывный выровненный буфер, так что даже спуск по $B$ превращается в последовательный стриминг линия за линией) и внутреннее регистровое микроядро, которое держит маленький блок $C$ целиком в векторных регистрах, пока через него идёт поток FMA. Результат умещается в один график:

OpenBLAS выдерживает 51–54 GFLOP/s от размера 256 до 4096 — плоская линия примерно на трёх четвертях от пунктирного пика одного ядра в 67 GFLOP/s — рис. 2 — однопоточный A @ B по размерам; плоская линия — и есть вся суть

Плоско. От $n = 256$ , где все три матрицы помещаются в L3, до $n = 4096$ , где 268 МБ операндов (400 МБ вместе с результатом) живут в DRAM, производительность держится на 51–54 GFLOP/s — 79% теоретического пика одного ядра на самом большом размере — потому что блочное разбиение делает отношение FLOP на байт свойством тайла, а не размера задачи. Цикл из определения тратит всё больше времени на каждый FLOP всякий раз, когда матрицы перерастают очередной уровень кэша; блочному циклу всё равно.

Две сноски с той же машины, обе измеренные. Первая — параллелизм: при привязке к четырём физическим ядрам $n = 4096$ выполняется на 179 GFLOP/s — ускорение в 3,4 раза; недобор до 4× — это в основном all-core turbo ниже одноядерного, плюс делёж L3 и каналов DRAM. Оставленный с настройками по умолчанию, OpenBLAS поднимает на этой машине восемь потоков — по одному на гиперпоток — и проседает до 118 GFLOP/s: два SMT-соседа делят FMA-конвейеры одного физического ядра, так что лишние потоки добавляют накладные расходы на планирование и ноль вычислений. Считайте физические ядра. Вторая — 79% пика: это уровень, на который зрелый BLAS выходит на этой микроархитектуре в float64; недостающая пятая часть уходит на трафик упаковки, края тайлов и накладные расходы на организацию циклов. Сто процентов не получает никто.

Та же идея — на всех этажах

Замените «L2» на «shared memory» — и это лекция по CUDA: GPU-ядро matmul складывает тайлы $A$ и $B$ в разделяемую память каждого SM, синхронизируется, умножает, сдвигается — тот же аргумент $\sqrt{M}$ , исполняемый тысячами потоков. Тензорные ядра переносят тайл в сам тракт данных: варп подаёт фрагменты фиксированного размера, и железо выполняет каждое маленькое плотное произведение одной инструкцией. Этажом выше слой внимания или MLP-блок — это батч таких произведений, и потому арифметическая интенсивность — теперь чаще в формулировке «FLOP на байт весов» — по-прежнему решает, будет ваш трансформер compute-bound или bandwidth-bound на H100.

Определение — это три цикла, по существу неизменные с тех пор, как Бине записал правило «строка на столбец» в 1812 году. Всё, что лежит между 9,3 секунды и 5,1 миллисекунды, — это знание о том, где ваши кэш-линии.

Воспроизвести

matmul_bench.py — ~130 строк Python со стандартной библиотекой и NumPy: два медленных цикла (по одному запуску), прогон BLAS по размерам матриц и демо обхода (лучший из 5), многопоточные замеры — всё с фиксированным сидом. Скрипт замеряет ровно те функции, что показаны выше. Абсолютные числа на вашей машине будут другими; соотношения — нет. Запустите, прежде чем верить мне на слово.