Что такое memory coalescing и почему оно важно для attention?

Краткий тезис

Memory coalescing — это механизм в GPU, при котором несколько потоков одного warp объединяют свои обращения к глобальной памяти в одну транзакцию, если они читают соседние адреса. Для attention (особенно в LLM) это критично, потому что при вычислении softmax(QK^T) и последующем чтении матриц K и V некоалесцированный доступ (например, при stride > 1) резко снижает пропускную способность памяти. Современные алгоритмы, такие как FlashAttention, используют tiling и работу в shared memory, чтобы гарантировать coalesced access и достичь высокой производительности.

1. Термин: Memory coalescing

Memory coalescing — это свойство архитектуры GPU (NVIDIA CUDA), при котором обращения нескольких потоков одного warp (группа из 32 потоков) к глобальной памяти объединяются в одну или несколько транзакций, если запрашиваемые адреса лежат в одном непрерывном блоке.

• Warp — минимальная единица исполнения в GPU: 32 потока выполняют одну инструкцию одновременно (SIMT).
• Транзакция памяти — минимальный объём данных, который GPU может прочитать из DRAM за один раз (обычно 32, 64 или 128 байт).
• Coalesced access — когда все 32 потока warp читают последовательные 32-, 64- или 128-байтовые слова, начиная с адреса, кратного размеру транзакции.

Если адреса разбросаны (stride > 1, случайный доступ), GPU вынуждена выполнять множество отдельных транзакций, что резко снижает эффективную пропускную способность памяти (bandwidth).

Пример:

• Coalesced: потоки 0..31 читают байты 0..31 → одна 32-байтовая транзакция.
• Strided: потоки читают байты 0, 128, 256, ... → 32 транзакции по 32 байта (или больше, в зависимости от шаблона).

2. Почему это важно для производительности GPU

GPU — это пропускная способность памяти (memory bandwidth) ограниченный процессор. Современные GPU имеют bandwidth 1–2 ТБ/с, но latency доступа к глобальной памяти составляет сотни тактов. Coalescing позволяет максимально использовать bandwidth: одна транзакция приносит данные для многих потоков, и warp не простаивает.

• Без coalescing: эффективная bandwidth падает в разы (например, 10% от пиковой).
• С coalescing: можно приблизиться к пиковой bandwidth (80–90%).

Для вычислительно-интенсивных операций (как attention) узким местом часто является чтение матриц K и V, поэтому coalescing — ключевой фактор производительности.

3. Как работает attention в LLM

Рассмотрим стандартный scaled dot-product attention:

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V

• Q, K, V — матрицы формы (batch_size, num_heads, seq_len, d_head).
• Вычисление S = Q * K^T — умножение матриц.
• Softmax по строкам S.
• Умножение результата на V.

На этапе чтения K и V из глобальной памяти (HBM) возникает проблема: при стандартной реализации каждый поток обрабатывает один элемент, и доступ к K/V может быть некоалесцированным.

4. Проблема разреженного доступа в attention

Рассмотрим типичную реализацию attention на GPU (без оптимизаций). Пусть:

• seq_len = N, d_head = D.
• Матрица K хранится в row-major: K[i][j] — элемент строки i, столбца j.
• При вычислении Q * K^T поток, отвечающий за элемент (i, j) матрицы S, читает строку Q[i] и столбец K[:, j].

Проблема: чтение столбца K[:, j] — это доступ с шагом D (stride = D). Потоки одного warp, обрабатывающие соседние j, читают адреса, отстоящие на D элементов. Если D > 1 (а обычно D = 64 или 128), то доступ некоалесцирован — каждый поток читает из разных строк, и транзакции не объединяются.

Аналогично для V при умножении softmax(S) * V: чтение V[:, j] также strided.

Результат: эффективная bandwidth падает в D раз (например, при D=128 — менее 1% пиковой).

5. Пример: naive attention vs coalesced attention

Ключевой приём: tiling — разбиение матриц на блоки, которые помещаются в shared memory (SRAM). Внутри блока доступ к данным уже coalesced, так как блок загружается непрерывно.

6. FlashAttention и tiling

FlashAttention (Dao et al., 2022) — алгоритм, который решает проблему некоалесцированного доступа и заодно уменьшает объём чтения/записи HBM.

Идея:

• Разбить Q, K, V на блоки (tiles) по seq_len.
• Загрузить блок K и V в shared memory (SRAM) — это coalesced access, так как читаются непрерывные строки.
• Вычислить частичное softmax и накопить результат в SRAM.
• Записать финальный результат обратно в HBM.

Почему это даёт coalescing:

• Загрузка K: каждый warp читает непрерывный фрагмент строк K (stride = 1) → coalesced.
• Загрузка V: аналогично.
• Внутри SRAM доступ быстрый и не требует coalescing (SRAM — это не глобальная память).

Дополнительно: FlashAttention использует online softmax (safe softmax) и recomputation градиентов для обратного прохода, что ещё больше снижает объём HBM-трафика.

7. Дополнительные техники: shared memory и bank conflicts

Даже при использовании shared memory могут возникнуть bank conflicts — когда несколько потоков обращаются к одному банку памяти, что снижает пропускную способность SRAM. Для их избегания применяют padding (добавление фиктивных столбцов) или перестановки.

Пример: если D = 64, а shared memory имеет 32 банка, то при stride = 2 (из-за padding) конфликтов не будет. FlashAttention учитывает это.

8. Связь с другими оптимизациями LLM inference

• PagedAttention (vLLM): использует страничную организацию KV cache, что также влияет на coalescing — при случайном доступе к страницам эффективность падает. Оптимизируется через prefetch и warp-level планирование.
• Tensor Parallelism: распределение голов attention по GPU — каждый GPU работает со своим блоком, внутри которого coalescing сохраняется.
• Quantization: снижение точности (FP16 → INT8) уменьшает объём данных, но не меняет паттерн доступа.

9. Итог: почему это критично для LLM inference

Attention — доминирующая операция в LLM (до 70% времени). Без coalescing эффективная bandwidth падает, и inference становится в разы медленнее. FlashAttention и его варианты (FlashAttention-2, FlashDecoding) — обязательные оптимизации для production-систем. Понимание memory coalescing необходимо для разработки эффективных kernel'ей и выбора аппаратного обеспечения.

Пет-проект для закрепления

Задача: Написать два простых CUDA kernel'а для attention (без softmax, только умножение Q*K^T) — один с strided доступом, другой с tiling и coalesced загрузкой. Сравнить время выполнения.

Инструменты: CUDA C++, nvcc, nvidia-smi, cudaEvent для замера времени.

Шаги:

1. Создать матрицы Q, K размером (N, D) в HBM (N=1024, D=64).
2. Kernel naive: каждый поток вычисляет один элемент S[i][j], читает Q[i][k] и K[k][j] (strided).
3. Kernel tiled: разбить K на блоки по 32 строки, загрузить блок в shared memory (coalesced), затем каждый поток читает из shared memory.
4. Замерить время для разных N (512, 1024, 2048).
5. Построить график ускорения.

Ожидаемый результат: tiled kernel в 5–10 раз быстрее для N=1024, D=64. При увеличении N разрыв растёт.

Связь с другими вопросами

Навигация

• Предыдущий: 701
• Следующий: 703
• Индекс: 00. Индекс разборов