ИИ и технологии

Самовнимание, ядра CUDA и FlashAttention

2026-11-15

Иллюстративное изображение: блок памяти микроконтроллера, не GPU, о котором идёт речь в статье.

Современные Transformer опираются на матричную математику самовнимания, высокоуровневые GPU-компиляторы вроде Triton и учитывающие память CUDA-ядра вроде FlashAttention, чтобы обрабатывать огромные последовательности токенов на пиковой скорости оборудования.

1. Как математически работает самовнимание в Vision Transformer

В Vision Transformer (ViT) изображение (например, 224×224 пикселей) нарезается на сетку патчей 16×16, что даёт N = 196 развёрнутых патчей. Каждый патч линейно проецируется в вектор размерности D (например, 768) и объединяется с позиционным эмбеддингом, образуя входную матрицу X ∈ R^(N × D).

Создание запросов, ключей и значений. Модель умножает матрицу патчей X на три обучаемые весовые матрицы (W_Q, W_K, W_V ∈ R^(D × d_k)), чтобы создать три разных представления для каждого патча изображения:

Q=XW_(Q), K=XW_(K), V=XW_(V)

Запрос (Q): представляет то, что ищет конкретный патч изображения (например, «Я текстура меха, ищу рядом глаза или уши»).
Ключ (K): объявляет, какую информацию содержит патч (например, «Я содержу резкий изогнутый край»).
Значение (V): хранит собственно полезную визуальную нагрузку признака, которая будет передана дальше в случае совпадения.

Вычисление оценок внимания. Чтобы измерить, насколько сильно каждый патч связан с каждым другим, модель вычисляет скалярное произведение всех запросов и ключей:

Scores=(QK^(T))/(sqrt(d_(k)))

Деление на sqrt(d_(k)) не даёт большим скалярным произведениям загнать следующую за ним функцию softmax в области с крошечными градиентами. В результате получается матрица внимания N × N, в которой элемент (i, j) оценивает, насколько патч i должен обращать внимание на патч j.

Взвешивание softmax и агрегирование. Применение softmax по строкам превращает необработанные оценки в вероятности, сумма которых равна 1. Умножение этих вероятностей на матрицу значений V даёт смешанное представление для каждого патча:

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

В многоголовом самовнимании тот же самый процесс выполняется параллельно в нескольких головах (например, 12 головах), позволяя разным головам одновременно сосредотачиваться на разных пространственных расстояниях, цветах или частях объектов.

2. Практическое сравнение кода: чистый CUDA C++ и OpenAI Triton

Чтобы увидеть, как эволюционировало GPU-программирование, рассмотрим простое поэлементное сложение векторов (C = A + B) для N элементов.

Чистый CUDA C++ (управление на уровне потоков)

В CUDA C++ вы программируете с точки зрения одного скалярного потока. Нужно вручную вычислять идентификаторы потоков, проверять границы массивов и писать шаблонный код конфигурации запуска на стороне хоста.

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton (программирование на Python на уровне блоков)

В Triton вы программируете прямо на Python с точки зрения блока (тайла) данных. Компилятор Triton автоматически берёт на себя синхронизацию потоков, объединение обращений к памяти (coalescing) и векторизацию «под капотом».

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention: решение проблемы узкого места памяти в Transformer

Стандартное самовнимание упирается в серьёзную стену при обработке изображений высокого разрешения или длинных текстовых последовательностей, поскольку матрица внимания S = QK^T растёт квадратично (N × N) с длиной последовательности N.

Узкое место пропускной способности памяти (HBM и SRAM). У современных GPU NVIDIA есть два основных уровня памяти:

HBM (высокоскоростная память / VRAM): большая (например, 80 ГБ), но относительно медленная при чтении и записи.
SRAM (встроенная разделяемая память): крошечная (около 20 МБ суммарно по всем потоковым мультипроцессорам), но примерно в 10–20 раз быстрее HBM.

В стандартном внимании GPU вычисляет S = QK^T, записывает огромную матрицу N × N в медленную HBM, считывает её обратно для применения softmax, снова записывает вероятности в HBM и считывает их ещё раз для умножения на V. Tensor Cores большую часть времени простаивают в ожидании передачи данных из памяти.

Тайлинг (никогда не материализовать матрицу N × N). Представленный Tri Dao и соавторами (2022), FlashAttention использует слияние ядер (kernel fusion), чтобы держать все вычисление внимания внутри быстрой встроенной SRAM. Он нарезает Q, K и V на небольшие блоки, помещающиеся в SRAM, вычисляет внимание тайл за тайлом и записывает обратно в HBM только итоговую выходную матрицу. Гигантская промежуточная матрица N × N никогда не сохраняется в VRAM.

Трюк с онлайновым softmax. Обычно softmax требует знать максимальное значение и сумму экспонент по всей строке, прежде чем нормализовать хотя бы один элемент. FlashAttention использует алгоритм онлайнового softmax, который блок за блоком отслеживает текущие максимумы и нормирующие множители, математически перемасштабируя частичные выходы по мере обработки каждого нового тайла.

Пересчёт при обратном распространении. Поскольку FlashAttention не сохраняет матрицу внимания N × N при прямом проходе, во время обучения он применяет умный компромисс: сохраняет в HBM только небольшую статистику нормализации и пересчитывает тайлы внимания на лету во внутренней быстрой SRAM при обратном распространении. Дополнительные вычисления в SRAM оказываются намного быстрее, чем чтение больших матриц из HBM, что даёт ускорение по реальному времени в 2–4 раза и сокращает расход памяти с квадратичного O(N²) до линейного O(N).

Больше новостей: ИИ и технологии