Как работает warp scheduling на NVIDIA GPU и как это влияет на LLM kernels?

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

Warp scheduling — ключевой механизм GPU, при котором Streaming Multiprocessor (SM) выбирает, какой warp (группа из 32 потоков) выполнять в данный момент. Для LLM-ядер (kernels), особенно attention, характерно множество ветвлений (if/else), что приводит к warp divergence — потоки внутри warp идут по разным путям, и GPU вынуждена сериализовать выполнение. Понимание warp scheduling помогает оптимизировать LLM-инференс: минимизировать divergence, повысить occupancy и скрыть latency за счёт переключения между warps.

1. Термины и базовая архитектура

1.1 Warp

Warp — минимальная единица исполнения на NVIDIA GPU: 32 потока, которые одновременно выполняют одну и ту же инструкцию (модель SIMT — Single Instruction, Multiple Threads). Если все 32 потока идут по одному пути, warp выполняется за один такт. Если пути расходятся — возникает divergence.

1.2 Streaming Multiprocessor (SM)

SM — вычислительный блок GPU, содержащий:

• несколько warp schedulers (например, 4 на H100);
• CUDA cores (FP32, INT32, Tensor Cores);
• разделяемую память, регистры, кэш L1.

Каждый SM может одновременно держать в состоянии (активных) до 64 warps (на H100 — до 64, на A100 — до 64, на V100 — до 64). Из них физически исполняться может только ограниченное число (зависит от числа schedulers и pipeline).

1.3 Warp scheduler

Warp scheduler — аппаратный блок SM, который на каждом такте выбирает один warp из пула готовых к исполнению (warp pool) и отправляет его инструкцию на исполнение. На H100 — 4 warp schedulers, каждый может выдавать по две инструкции за такт (dual-issue).

1.4 Warp divergence

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

1.5 Latency hiding

Latency hiding — способность GPU скрывать задержки (например, доступ к глобальной памяти) путём переключения на другой warp, пока текущий ждёт данные. Чем больше активных warps, тем лучше скрывается latency.

2. Механизм warp scheduling

2.1 Пул готовых warps

На каждом SM есть warp pool — все warps, назначенные этому SM (до 64). Из них часть может быть в состоянии ready (все операнды готовы, инструкция может быть выполнена), часть — stalled (ожидание памяти, синхронизации и т.д.).

2.2 Алгоритм выбора

Scheduler использует простой алгоритм: на каждом такте он выбирает один из ready warps. Обычно применяется round-robin с приоритетом по возрасту (oldest-ready-first). На современных GPU (Volta+) используется two-level scheduler:

• Level 1: выбор warp из набора активных warps (по приоритету).
• Level 2: выбор инструкции внутри warp (dual-issue).

2.3 Пример: H100

• 4 warp schedulers на SM.
• Каждый scheduler может выдать до 2 инструкций за такт (одна для CUDA core, одна для Tensor Core или специальных блоков).
• Максимальная пропускная способность: 4 × 2 = 8 инструкций за такт на SM.
• Чтобы полностью занять schedulers, нужно минимум 8 ready warps (по одному на каждый issue slot). На практике нужно больше для скрытия latency.

2.4 Влияние occupancy

Occupancy — отношение числа активных warps к максимально возможному. Высокий occupancy (например, 64 warps) даёт больше возможностей для переключения, но не гарантирует высокой производительности, если warps часто stalled из-за divergence или неэффективного доступа к памяти.

3. Warp divergence в LLM kernels

3.1 Почему LLM kernels склонны к divergence

LLM-ядра, особенно attention, содержат много условных операций:

• Causal mask: для каждого query нужно маскировать токены будущего (if (position > query_position) then set to -inf).
• Variable sequence lengths: в batch-обработке длины последовательностей разные → padding и маски.
• Softmax: вычисление экспоненты и суммы — нет ветвлений, но есть деление на ноль (редко).
• FlashAttention: использует tiling и recomputation, но внутри tile тоже может быть маскирование.

3.2 Пример divergence в attention

// Псевдокод attention kernel
for (int i = 0; i < seq_len; ++i) {
    float score = dot(query, key[i]);
    if (i > query_idx) {  // causal mask
        score = -INFINITY;
    }
    // softmax accumulation
}

Здесь if (i > query_idx) — ветвление, зависящее от query_idx. Внутри warp потоки могут иметь разные query_idx (разные запросы), поэтому одни потоки выполнят score = -INFINITY, другие — нет. GPU выполнит обе ветви последовательно.

3.3 Последствия

• Serialization: вместо одной инструкции — две (или больше).
• Снижение occupancy: warps с divergence дольше заняты, меньше готовых warps для переключения.
• Падение utilisation Tensor Cores: Tensor Cores требуют uniform data layout; divergence может вынудить использовать CUDA cores вместо Tensor Cores.

4. FlashAttention и минимизация divergence

4.1 Uniform control flow

FlashAttention (Dao et al., 2022) перепроектирует attention так, чтобы все потоки внутри warp выполняли одинаковые операции:

• Tiling: блоки Q, K, V фиксированного размера (например, 32×32).
• Online softmax: вычисляется по tile, без глобальной синхронизации.
• Recomputation: не хранить большие матрицы attention, а пересчитывать на лету.

В результате внутри warp нет ветвлений, зависящих от данных: все потоки обрабатывают один и тот же tile. Маскирование реализуется через predication (условное присваивание), которое не вызывает divergence.

4.2 Predication vs divergence

Predication — техника, при которой обе ветви выполняются, но результаты одной маскируются (записываются в регистр, но не сохраняются). На GPU это часто эффективнее divergence, если ветви короткие. FlashAttention использует predication для causal mask.

4.3 Производительность

FlashAttention достигает 2–3× ускорения по сравнению с стандартным attention на длинных последовательностях, во многом благодаря устранению divergence и эффективному использованию памяти.

5. Скрытие latency переключением warps

5.1 Формула скрытия latency

Для полного скрытия latency доступа к глобальной памяти (например, 200–400 тактов) нужно достаточное количество warps:

Warps_needed = (memory_latency_cycles) / (issue_interval_cycles)

На H100: memory latency ~300 тактов, issue interval (время между инструкциями одного warp) ~4 такта → нужно ~75 warps. Но SM может держать только 64 warps, поэтому часть latency не скрывается. Однако для LLM kernels часто bottleneck — compute, а не memory, поэтому occupancy не критична.

5.2 Роль warp scheduling

Warp scheduler автоматически переключается на ready warp, когда текущий stall. Если divergence уменьшает число ready warps, latency скрывается хуже. Поэтому для LLM kernels важно поддерживать высокий occupancy и минимизировать stall-факторы (divergence, bank conflicts).

6. Практические рекомендации для LLM kernels

6.1 Инструменты профилирования

• Nsight Compute: показывает warp stall reasons, divergence rate, occupancy.
• NVIDIA Nsight Systems: временные диаграммы выполнения warps.
• CUDA occupancy calculator: расчёт максимального числа warps при заданном использовании ресурсов.

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

Задача: Написать простой CUDA kernel для causal attention, измерить warp divergence и сравнить с версией, использующей predication.

Инструменты: CUDA Toolkit, Nsight Compute, Python (для запуска и анализа).

Шаги:

1. Реализовать kernel attention_divergent: для каждого query вычислять dot product со всеми key, использовать if (i > query_idx) score = -INFINITY.
2. Реализовать kernel attention_predicated: вместо if использовать score = (i <= query_idx) ? dot(...) : -INFINITY; (тернарный оператор компилируется в predication).
3. Запустить оба kernel на случайных данных (batch=1, seq_len=1024, head_dim=64).
4. Использовать Nsight Compute для профилирования: сравнить divergence branches, stall reasons, occupancy.
5. Измерить время выполнения (CUDA events).

Ожидаемый результат: predicated kernel покажет меньше divergence (или 0), меньше stall из-за divergence, выше производительность (на 10–30% в зависимости от длины).

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

Навигация

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