Проблема: Attention доминирует во времени вывода
По мере роста контекстных окон (до 128K+ токенов) операция Attention становится главным узким местом. На примере DeepSeek-R1 показано, что доля времени, затрачиваемого на Attention, возрастает с 18% при контексте 4K до 85% при 128K. Это требует пересмотра архитектуры моделей под конкретные GPU-архитектуры NVIDIA, а не просто оптимизации кода.
Ключевые метрики: Prefill vs Decode
Фазы prefill (обработка промпта) и decode (генерация токенов) имеют разные ограничения. Prefill ограничен вычислительной мощностью (compute-bound), так как оперирует большими матрицами. Decode ограничен пропускной способностью памяти (memory-bound), так как читает KV-cache из HBM для каждого нового токена.
| Параметр | Prefill | Decode |
|---|---|---|
| Длина запроса | Полный ввод (ISL токенов) | 1 токен |
| Контекст Attention | Промпт (ISL токенов) | Весь KV-cache (KVSL токенов) |
| Основное ограничение | Вычисления (matmul + softmax) | Пропускная способность HBM |
| Размер GEMM-M | ISL (большой) | 1 (маленький) |
Решение 1: Group Size (G) критичен для Decode
Group size ($G$) — отношение количества query-голов к KV-головам. Для фазы decode увеличение $G$ напрямую повышает арифметическую интенсивность и снижает объем передаваемых данных. Увеличение $G$ с 1 (MHA) до 8 дает ускорение в 2 раза, так как уменьшается объем чтения из памяти. NVIDIA рекомендует использовать GQA с высоким $G$ (например, Nemotron 3 использует всего 2 KV-головы) для максимальной эффективности decode.
Решение 2: Head Dimension и Parallelism
Для совпадения с тайлингом GPU и выравниванием памяти оптимальными размерами head dimension ($Hsz$) являются 128 или 256. Стратегия параллелизма зависит от количества KV-голов ($KH$):
- Tensor Parallelism (TP): Не должно превышать количество KV-голов, чтобы избежать дублирования KV-состояния.
- Малое количество KV-голов: Требует использования Attention Data Parallelism (ADP), KV Parallelism (KVP) или гибридных подходов (Wide EP, Helix Parallelism), реализованных в TensorRT-LLM.
Практические выводы
Для достижения максимальной пропускной способности на NVIDIA GPU разработчикам следует:
- Максимизировать group size ($G$) для фазы decode.
- Использовать head dimensions 128 или 256.
- Применять сжатие KV-кэша, sparse-attention или sliding-window для минимизации эффективного состояния KV.
- Адаптировать стратегию параллелизма под количество KV-голов модели.
Источник: NVIDIA dev blog ↗