Grouped Query Attention (GQA)

Grouped Query Attention (GQA) — это архитектура механизма внимания, в которой несколько голов запросов делят между собой меньший набор голов ключей и значений. Она сохраняет разнообразие точек зрения многоголового внимания (multi-head attention), уменьшая при этом объем памяти и перемещение данных, необходимые для ключей и значений. GQA особенно ценна при авторегрессионной генерации, когда модель создает по одному токену за раз и многократно считывает сохраненные состояния внимания.

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

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

• Query: описывает информацию, которую ищет текущий токен.
• Key: описывает, что представляет собой каждый доступный токен.
• Value: содержит информацию, которая извлекается, когда query совпадает с key.

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

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

GQA в сравнении с Multi-Head и Multi-Query Attention#

GQA занимает промежуточное положение между многоголовым вниманием (multi-head attention) и вниманием с одним запросом (multi-query attention):

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

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

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

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

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

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

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

Реальные приложения#

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

Мультимодальные ассистенты по работе с изображениями и видео: vision-language model может представлять патчи изображений или видеокадры в виде длинных последовательностей токенов. 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 позволяет управлять тем, какое поддерживаемое ядро выполняет операцию, а руководство по строительным блокам Transformer в PyTorch дает более широкий контекст реализации. (docs.pytorch.org)

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