Как работает speculative execution на GPU для LLM (branch prediction)?

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

GPU не использует сложное branch prediction (предсказание ветвлений) как CPU из-за архитектуры SIMT (Single Instruction, Multiple Threads). Вместо этого GPU применяет predicated execution (исполнение с предикатами]]): при ветвлении все потоки warp (группы из 32 потоков) выполняют обе ветки, а результаты маскируются. Для LLM-ядер критично минимизировать divergent control flow (расходящиеся ветвления) и использовать predicated инструкции (selp, setp). Пример — FlashAttention, где внутри блока потоков достигается uniform control flow (единообразный поток управления).

1. Термины и контекст

Speculative execution (execution|спекулятивное исполнение) — техника CPU, при которой процессор предсказывает исход условного перехода и начинает выполнять одну из веток до проверки условия. При ошибке предсказания результаты отбрасываются, что вызывает штраф (pipeline flush).

prediction|Branch prediction (предсказание ветвлений) — компонент CPU, предсказывающий, какая ветка будет выполнена. Современные CPU достигают точности >95%, но ошибки всё равно дороги.

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

execution|Predicated execution (исполнение с предикатами) — механизм GPU: вместо предсказания ветки, GPU выполняет обе ветки, но только для тех потоков, где условие истинно. Результаты для неактивных потоков маскируются (не записываются).

control control flow|Divergent control control flow (control control расходящееся управление) — ситуация, когда внутри warp разные потоки выбирают разные ветки условного оператора. Это снижает эффективность, так как все потоки выполняют обе ветки последовательно.

Uniform control flow (единообразное управление) — все потоки warp следуют одному пути, ветвления нет. Это оптимально для GPU.

2. Отличие CPU branch prediction от GPU predicated execution

Ключевой вывод: GPU не делает speculative execution в классическом смысле. Вместо этого он жертвует производительностью части потоков ради простоты и масштабируемости.

3. Как GPU обрабатывает ветвления: детальный механизм

Рассмотрим код на CUDA:

__global__ void kernel(float* data, int n) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < n) {
        if (data[idx] > 0.0f) {
            data[idx] = data[idx] * 2.0f;
        } else {
            data[idx] = data[idx] + 1.0f;
        }
    }
}

Шаги исполнения внутри warp

1. Проверка условия: инструкция SETP (set predicate) вычисляет предикат p = (data[idx] > 0.0f) для каждого потока.
2. Маскирование: GPU создаёт две маски — активные потоки для true-ветки (p == true) и для false-ветки (p == false).
3. Выполнение true-ветки: все потоки warp выполняют инструкцию умножения, но только те, у кого p == true, записывают результат. Остальные игнорируют запись (маскированы).
4. Выполнение false-ветки: аналогично, все потоки выполняют инструкцию сложения, записывают только те, у кого p == false.
5. Синхронизация: warp завершает обе ветки, и все потоки продолжают с unified кода.

Проблема: если внутри warp есть дивергенция, GPU тратит время на выполнение обеих веток. Если одна ветка короткая, а другая длинная, эффективность падает до 50% (половина потоков простаивает).

Инструкция selp (select with predicate) — оптимизированная версия: вычисляет оба результата и выбирает нужный по предикату без явного ветвления. Пример:

float val = data[idx];
float true_res = val * 2.0f;
float false_res = val + 1.0f;
data[idx] = (val > 0.0f) ? true_res : false_res; // компилятор может заменить на selp

4. Влияние на LLM kernels

LLM-ядра (например, attention, feed-forward) часто содержат условные операции:

• Маски внимания (causal mask, padding mask) — для каждого токена нужно решить, учитывать ли другие токены.
• Softmax — вычисление экспонент и суммирование, где есть проверки на -inf.
• KV cache — управление памятью с условными копированиями.
• Activation functions (ReLU, GELU) — могут иметь ветвления (например, if (x > 0)).

Пример дивергенции в attention:

for (int j = 0; j < seq_len; ++j) {
    if (mask[idx * seq_len + j] > 0) { // padding
        float score = query * key[j];
        // ...
    }
}

Здесь разные потоки warp могут иметь разные маски → дивергенция.

Последствия: снижение occupancy (загрузки GPU) и пропускной способности. Для LLM инференса это критично, так как latency чувствителен.

5. Оптимизации для LLM на GPU

5.1 Использование predicated инструкций

Компилятор NVIDIA (nvcc) автоматически заменяет простые if-else на selp, если обе ветки короткие. Разработчик может помочь:

• Избегать глубоких вложенных ветвлений.
• Использовать тернарный оператор ? : вместо if-else.
• Применять встроенные функции __ballot_sync, __any_sync для проверки условий внутри warp.

5.2 Минимизация дивергенции

• Реструктуризация данных: группировать потоки так, чтобы внутри warp все шли по одному пути. Например, сортировать токены по длине маски (bucketing).
• Uniform branching: если условие одинаково для всех потоков warp (например, if (blockIdx.x == 0)), дивергенции нет — это uniform control flow.
• Loop unrolling: разворачивать циклы с известным числом итераций, чтобы компилятор мог оптимизировать.

5.3 Пример: FlashAttention

FlashAttention — алгоритм attention, который разбивает вычисления на блоки и использует tiling. Внутри каждого блока все потоки warp выполняют одинаковые операции: загрузка данных, умножение, softmax. Условия (например, проверка на конец последовательности) вынесены на уровень блоков, а не потоков. Это даёт uniform control flow и высокую эффективность.

// FlashAttention kernel — нет if внутри warp
for (int j = 0; j < num_blocks; ++j) {
    float4 key_block = load_key_block(j);
    float score = dot(query, key_block);
    // softmax — uniform операции
}

6. Speculative execution в контексте LLM: не путать с speculative decoding

Speculative execution на GPU — это не то же самое, что speculative decoding (спекулятивная декодировка) для LLM. Speculative decoding — техника ускорения генерации текста, при которой маленькая модель (draft model) предсказывает несколько токенов, а большая модель проверяет их. Это не связано с branch prediction.

В контексте GPU speculative execution может относиться к warp scheduling — когда один warp простаивает (ждёт память), GPU переключается на другой warp. Это скрывает задержки, но не является предсказанием ветвлений.

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

NVIDIA Nsight Compute — профилировщик CUDA ядер. Позволяет увидеть:

• Divergence — процент дивергентных ветвлений в warp.
• Branch efficiency — доля потоков, которые реально выполнили инструкцию (не маскированы).
• Occupancy — загрузка SM (Streaming Multiprocessor).

Команда для профилирования:

ncu --set full -o profile ./my_llm_kernel

В GUI можно найти раздел Warp State Statistics и Branch Statistics.

8. Практические рекомендации для разработчиков LLM kernels

1. Измеряйте дивергенцию — используйте Nsight Compute. Если divergence > 10%, ищите оптимизации.
2. Используйте __syncthreads() — синхронизация внутри блока может помочь унифицировать управление.
3. Предпочитайте selp — компилятор часто сам заменяет, но можно явно написать через ? :.
4. Группируйте данные — сортируйте последовательности по длине (bucketing) для uniform масок.
5. Избегайте ветвлений в горячих циклах — выносите условия наружу.
6. Используйте #pragma unroll — помогает компилятору оптимизировать.

9. Связь с другими аспектами GPU

• Tensor Cores — специализированные блоки для матричных умножений. Они не поддерживают ветвления, поэтому код должен быть uniform.
• Memory coalescing — дивергенция может нарушить объединение запросов к памяти, снижая пропускную способность.
• Occupancy — дивергенция снижает количество активных warp на SM, уменьшая occupancy.

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

Задача: Написать простой CUDA kernel для фрагмента attention с дивергентным ветвлением и оптимизировать его.

Инструменты: CUDA Toolkit, Nsight Compute, Python (для тестов).

Шаги:

1. Реализовать kernel, который для каждого потока вычисляет score = dot(query, key), и если score > threshold, то добавляет его в сумму (ветвление).
2. Измерить производительность и дивергенцию через ncu.
3. Оптимизировать: заменить if на selp (через тернарный оператор), переписать без ветвления, используя маску.
4. Сравнить время выполнения и branch efficiency.

Ожидаемый результат: Понимание, как дивергенция влияет на производительность, и умение применять predicated execution.

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

Навигация

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