Grouped Query Attention (GQA)

Сгруппированное внимание по запросам (GQA) — это архитектура внимания, в которой несколько голов запросов используют общий меньший набор голов ключей и значений. Она сохраняет разнообразие представлений многоголовых запросов, одновременно сокращая объем памяти и передачи данных, необходимых для ключей и значений. GQA особенно полезно при авторегрессионном выводе, когда модель генерирует по одному токену за раз и многократно считывает сохраненные состояния внимания.

Как работает GQA#

В рамках механизма внимания каждый токен проецируется в три представления:

• Запрос: описывает информацию, которую ищет текущий токен.
• Ключ: описывает, что представляет каждый доступный токен.
• Значение: содержит информацию, извлекаемую при совпадении запроса с ключом.

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

Сама формула внимания остается вниманием на основе масштабированного скалярного произведения. Меняются расположение голов и объем данных ключей и значений, которые создаются, сохраняются и считываются. Поэтому такие фреймворки, как масштабированное скалярное произведение внимания в PyTorch, требуют, чтобы при включенном GQA число голов запросов делилось на число голов ключей и значений без остатка. (docs.pytorch.org)

GQA по сравнению с многоголовым вниманием и вниманием с несколькими запросами#

GQA занимает промежуточное положение между многоголовым вниманием и вниманием с несколькими запросами:

• Многоголовое внимание: у каждой головы запросов есть собственные головы ключей и значений. Это обеспечивает максимальную независимость голов, но требует наибольшего объема памяти для ключей и значений.
• Сгруппированное внимание по запросам: несколько голов запросов используют каждую голову ключей и значений совместно. Это обеспечивает баланс между емкостью представления и эффективностью использования памяти.
• Внимание с несколькими запросами: все головы запросов используют одну голову ключей и одну голову значений. Это минимизирует объем хранилища ключей и значений, но обеспечивает меньшее разнообразие ключей и значений.

Если слой имеет 32 головы запросов, многоголовое внимание также может использовать 32 головы ключей и значений, GQA — 8, а внимание с несколькими запросами — 1. GQA не заменяет более широкую архитектуру Transformer; это один из возможных вариантов конфигурации внутри слоя внимания Transformer. В документации NVIDIA по многоголовому вниманию, вниманию с несколькими запросами и сгруппированному вниманию по запросам описывается, что эти варианты различаются прежде всего количеством голов ключей и значений, обслуживающих головы запросов. (nvidia.github.io)

Почему GQA имеет значение#

При авторегрессионной генерации модель сохраняет предыдущие ключи и значения в кэше KV. При стандартном многоголовом внимании каждая голова внимания добавляет отдельные кэшированные ключи и значения. Поэтому длинные последовательности, большие пакеты и множество слоев модели могут потреблять значительный объем памяти GPU.

Поскольку GQA использует меньше голов ключей и значений, его кэш уменьшается пропорционально сокращению числа этих голов. Это может улучшить:

• Емкость памяти: в доступной памяти помещаются более длинные запросы или большее количество одновременных запросов.
• Пропускную способность генерации: для каждого сгенерированного токена требуется считывать меньше данных ключей и значений.
• Задержка вывода: сокращение трафика памяти может уменьшить время обработки между токенами.
• Масштабируемость контекстного окна: обслуживание длинных последовательностей становится более практичным.

Эти преимущества зависят от аппаратного обеспечения, длины последовательности, размера пакета и поддержки ядер. Средам выполнения в production по-прежнему требуется эффективное выделение кэша; например, система кэша KV в TensorRT-LLM от NVIDIA объединяет поддержку GQA с блочным кэшированием, повторным использованием и выгрузкой. (nvidia.github.io)

Практические применения#

Ассистенты для работы с длинным контекстом: ассистент для программирования может обработать тысячи токенов исходного кода перед генерацией ответа. GQA уменьшает объем кэшированного состояния ключей и значений, связанного с этой историей, позволяя серверу обрабатывать более длинные файлы или большее количество одновременных пользователей без пропорционального увеличения памяти GPU.

Мультимодальные ассистенты для работы с изображениями и видео: модель компьютерного зрения и языка может представлять фрагменты изображений или кадры видео в виде длинных последовательностей токенов. GQA снижает нагрузку на кэш внимания после передачи этих визуальных токенов в языковой декодер, потенциально оставляя больше памяти для дополнительных изображений, кадров или запросов. Это отличается от стандартного Vision Transformer, в котором внимание может обрабатывать все фрагменты изображения параллельно, без авторегрессионного кэширования KV.

Практическая реализация и компромиссы#

PyTorch предоставляет доступ к GQA через enable_gqa=True. В этом документированном примере для CUDA используются восемь голов запросов и две общие головы ключей и значений:

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

assert torch.cuda.is_available(), "A CUDA device is required."

query = torch.randn(1, 8, 16, 32, device="cuda")
key = torch.randn(1, 2, 16, 32, device="cuda")
value = torch.randn(1, 2, 16, 32, device="cuda")

# Four query heads share each key-value head.
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        enable_gqa=True,
    )

print(output.shape)

На выходе сохраняются восемь голов запросов, тогда как для ключей и значений используются только две головы. Селектор бэкенда внимания PyTorch может управлять тем, какое поддерживаемое ядро выполняет операцию, а руководство PyTorch по строительным блокам Transformer содержит более широкий контекст реализации. (docs.pytorch.org)

GQA необходимо выбирать при проектировании или адаптации модели; обычно это не переключатель времени вывода для уже существующего чекпойнта. Разработчикам следует проверить поддержку фреймворка, делимость числа голов, численное поведение, использование памяти и качество выполнения задачи. Для развертывания компьютерного зрения дополнительные рабочие процессы, такие как интеграция Ultralytics с TensorRT и режим бенчмаркинга Ultralytics, помогают измерить, дают ли архитектурные оптимизации и оптимизации среды выполнения существенные улучшения на целевом оборудовании.