Что такое warp divergence в CUDA и как он влияет на attention?

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

Warp divergence — это ситуация, когда потоки внутри одного warp (группа из 32 потоков в CUDA) выполняют разные инструкции из-за условных переходов (if/else). Это приводит к сериализации выполнения: GPU выполняет все ветки последовательно, маскируя неактивные потоки, что резко снижает утилизацию вычислительных блоков. В контексте attention (особенно в трансформерах) warp divergence возникает из-за padded sequences (последовательностей разной длины, дополненных до максимума) и masking|causal masking, когда одни потоки обрабатывают реальные токены, а другие — паддинг. Это замедляет ядра attention, и для борьбы с этим применяются техники вроде Flash Attention и tiling.

1. Термин: Warp и SIMT-модель

Warp — это базовая единица планирования в CUDA. GPU NVIDIA объединяет 32 потока в один warp, и все они выполняют одну и ту же инструкцию одновременно (модель SIMT — Single Instruction, Multiple Threads). Потоки в warp имеют свои регистры и могут обращаться к разным данным, но на каждом такте все активные потоки исполняют одну и ту же команду из program counter.

SIMT отличается от SIMD (Single Instruction, Multiple Data) тем, что потоки могут независимо ветвиться, но с оговорками — если ветвление разное, возникает divergence.

2. Что такое warp divergence

Warp divergence происходит, когда внутри одного warp разные потоки выбирают разные пути исполнения (например, одни идут в if, другие — в else). Аппаратура не может выполнять разные инструкции одновременно, поэтому она сериализует ветки: сначала выполняет все потоки, вошедшие в if (остальные маскируются), затем — все потоки, вошедшие в else. Пока одна ветка активна, другая простаивает.

Пример кода с divergence

__global__ void divergent_kernel(float* data, int n) {
    int idx = threadIdx.x;
    if (idx < n / 2) {
        data[idx] = data[idx] * 2.0f;  // ветка A
    } else {
        data[idx] = data[idx] + 1.0f;  // ветка B
    }
}

Если warp состоит из 32 потоков, а n/2 = 16, то первые 16 потоков идут в ветку A, вторые 16 — в ветку B. GPU выполнит обе ветки последовательно, утилизация — 50%.

Когда divergence не происходит

• Если все потоки в warp принимают одинаковое решение (все идут в if или все в else) — divergence нет.
• Если условие зависит от uniform значения (одинакового для всех потоков), компилятор может оптимизировать.

3. Как divergence влияет на производительность

Основные последствия:

• Сериализация: общее время выполнения = сумма времени всех веток.
• Снижение occupancy (занятости SM): пока warp выполняет одну ветку, другая часть потоков простаивает, но warp всё равно занимает ресурсы (регистры, shared memory).
• Увеличение числа инструкций: каждая ветка добавляет дополнительные команды (предикаты, маски).

Количественная оценка Если в warp две ветки равной длины, производительность падает в 2 раза. Если ветки сильно разбалансированы (одна короткая, другая длинная), то проигрыш меньше, но всё равно есть.

4. Warp divergence в attention-механизме

Attention (особенно scaled dot-product attention) в трансформерах вычисляется как:

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

В реализации на GPU каждый поток или блок потоков обрабатывает один или несколько элементов (например, один query, или одну позицию в последовательности).

Основные источники divergence в attention

4.1 Padded sequences (последовательности разной длины)

В батче обычно все последовательности дополняются до максимальной длины (padding). При вычислении attention для коротких последовательностей часть токенов — паддинг (обычно маскируется). Если warp обрабатывает несколько query из разных последовательностей, то для одних query маска активна (нужно считать attention только до реальной длины), для других — нет. Это приводит к divergence: одни потоки выполняют вычисления для реальных токенов, другие — пропускают или обрабатывают паддинг.

4.2 Causal masking (авторегрессивное внимание)

В decoder-части трансформера используется causal mask, запрещающий токенам "смотреть" в будущее. Для каждого query маска разная: первый токен видит только себя, второй — себя и первый и т.д. Если warp обрабатывает несколько query подряд (например, в ядре attention с параллелизмом по query), то для каждого query количество активных ключей разное. Это вызывает divergence, так как потоки выполняют разное количество итераций цикла по ключам.

4.3 Variable-length sequences внутри батча

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

5. Пример: naive CUDA-ядро attention с divergence

Рассмотрим упрощённое ядро attention, где каждый поток обрабатывает один query (batch_size=1, но несколько потоков на разные query).

__global__ void attention_kernel(float* Q, float* K, float* V, float* out,
                                 int N, int d) {
    int q_idx = blockIdx.x * blockDim.x + threadIdx.x; // query index
    if (q_idx >= N) return;

    float sum = 0.0f;
    float max_val = -INFINITY;
    // loop over keys
    for (int k_idx = 0; k_idx < N; ++k_idx) {
        // causal mask: if k_idx > q_idx, skip
        if (k_idx <= q_idx) {
            float score = dot_product(Q + q_idx*d, K + k_idx*d, d);
            max_val = fmaxf(max_val, score);
        }
    }
    // ... softmax and weighted sum
}

В этом коде для каждого q_idx количество итераций цикла разное (от 1 до N). Если warp содержит потоки с разными q_idx, то условие k_idx <= q_idx будет истинно для разного числа итераций. Это приводит к divergence на каждой итерации цикла.

6. Как измерять warp divergence

Инструменты:

• NVIDIA Nsight Compute — показывает branch efficiency (доля инструкций, где все потоки в warp выбрали одну ветку). Низкий branch efficiency (< 90%) указывает на divergence.
• CUDA occupancy calculator — оценивает влияние divergence на occupancy.
• Профилирование: счётчики branch_efficiency, divergent_branch.

Типичные значения

• В хорошо оптимизированных ядрах attention (Flash Attention) branch efficiency > 95%.
• В наивных реализациях с padding может падать до 50-70%.

7. Методы борьбы с warp divergence в attention

7.1 Flash Attention

Flash Attention (Dao et al., 2022) использует tiling — разбиение матриц Q, K, V на блоки, которые помещаются в shared memory. Внутри блока все потоки обрабатывают одинаковое количество элементов, divergence минимизируется. Кроме того, Flash Attention не хранит полную матрицу attention, а вычисляет softmax по частям, что также снижает расхождение.

7.2 Padding-free батчинг

Группировка последовательностей одинаковой длины в батч (bucketing) или динамическое выравнивание. Например, xFormers и vLLM используют paged attention, где последовательности разбиваются на страницы, и warp всегда обрабатывает фиксированное количество токенов.

7.3 Использование предикатов вместо ветвления

В CUDA можно заменить if на предикатные инструкции (через __ballot_sync, __any_sync), но это не всегда эффективно. Компилятор часто сам преобразует короткие ветки в предикаты.

7.4 Перестановка данных (data reordering)

Переупорядочивание потоков так, чтобы в одном warp оказались query с одинаковой маской (например, все query из коротких последовательностей). Это требует дополнительной сортировки, но может повысить branch efficiency.

8. Сравнение: naive vs оптимизированное внимание

9. Практический пример: влияние padding на divergence

Допустим, батч из 4 последовательностей длин: [4, 8, 12, 16]. Максимальная длина = 16. Если warp обрабатывает 32 query (по 8 на каждую последовательность), то для первых 8 query (длина 4) маска активна только для 4 ключей, для следующих 8 — для 8 и т.д. Внутри warp потоки имеют разное количество итераций → divergence.

Решение сгруппировать последовательности по длине: батч из 4 последовательностей длины 4, батч из 4 длины 8 и т.д. Тогда внутри каждого батча все warp имеют одинаковую маску, divergence отсутствует.

10. Связь с другими оптимизациями GPU

Warp divergence — лишь один из факторов производительности CUDA-ядер. Другие важные аспекты:

• Memory coalescing — объединение обращений к глобальной памяти.
• Occupancy — количество активных warp на SM.
• Bank conflicts в shared memory.
• Instruction-level parallelism.

В attention критичны все эти аспекты, но divergence особенно важен из-за маскирования.

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

Задача Написать два CUDA-ядра для вычисления attention (causal mask) — наивное и оптимизированное с устранением divergence — и сравнить их производительность.

Инструменты

• CUDA Toolkit (nvcc)
• Python с PyTorch (для генерации данных и вызова ядер через torch.cuda)
• Nsight Compute для профилирования

Шаги:

1. Реализовать наивное ядро, где каждый поток обрабатывает один query, и цикл по ключам с условием if (k_idx <= q_idx).
2. Реализовать оптимизированное ядро, где warp обрабатывает блок query, и все потоки в warp имеют одинаковую маску (например, сгруппировать query по длине).
3. Запустить оба ядра на наборе последовательностей разной длины (например, от 64 до 1024 токенов) и измерить время выполнения.
4. Использовать Nsight Compute для сравнения branch efficiency.

Ожидаемый результат

• Наивное ядро покажет branch efficiency ~60-80% и время, растущее с разбросом длин.
• Оптимизированное ядро — branch efficiency >95% и ускорение в 1.5–3 раза.

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

Навигация

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