Grouped Query Attention (GQA)

Grouped Query Attention (GQA) is an attention design in which multiple query heads share a smaller set of key and value heads. It preserves the diverse perspectives of multi-head queries while reducing the memory and data movement required for keys and values. GQA is especially valuable during autoregressive inference, where a model generates one token at a time and repeatedly reads stored attention states.

How GQA Works#

Within an attention mechanism, each token is projected into three representations:

• Query: Describes the information the current token is seeking.
• Key: Describes what each available token represents.
• Value: Contains the information retrieved when a query matches a key.

In conventional self-attention, these projections are divided into heads so the model can learn different relationships in parallel. GQA keeps many query heads but assigns groups of them to shared key and value heads. For example, an attention layer may use eight query heads and two key-value heads. Each key-value head then serves four query heads.

The attention calculation itself remains scaled dot-product attention. What changes is the head arrangement and the amount of key-value data produced, stored, and read. Frameworks such as PyTorch scaled dot-product attention therefore require the number of query heads to be divisible by the number of key-value heads when GQA is enabled. (docs.pytorch.org)

GQA vs. Multi-Head and Multi-Query Attention#

GQA sits between multi-head attention and multi-query attention:

• Multi-head attention: Every query head has its own key and value heads. This offers maximum head independence but creates the largest key-value memory requirement.
• Grouped query attention: Several query heads share each key-value head. It balances representational capacity with memory efficiency.
• Multi-query attention: All query heads share one key head and one value head. This minimizes key-value storage but provides less key-value diversity.

If a layer has 32 query heads, multi-head attention may also use 32 key-value heads, GQA might use 8, and multi-query attention uses 1. GQA is not a replacement for the broader Transformer architecture; it is one possible configuration inside a Transformer attention layer. NVIDIA’s multi-head, multi-query, and grouped-query attention documentation describes these variants as differing primarily in how many key-value heads serve the query heads. (nvidia.github.io)

Why GQA Matters#

During autoregressive generation, a model stores previous keys and values in a KV cache. With standard multi-head attention, every attention head contributes separate cached keys and values. Long sequences, large batches, and many model layers can therefore consume substantial GPU memory.

Because GQA uses fewer key-value heads, its cache is smaller in proportion to the reduction in those heads. This can improve:

• Memory capacity: Longer prompts or more simultaneous requests can fit in available memory.
• Generation throughput: Less key-value data must be read for every generated token.
• Inference latency: Reduced memory traffic can shorten token-to-token processing time.
• Context window scalability: Serving long sequences becomes more practical.

These gains depend on hardware, sequence length, batch size, and kernel support. Production runtimes still need efficient cache allocation; NVIDIA’s TensorRT-LLM KV cache system, for example, combines GQA support with block-based caching, reuse, and offloading. (nvidia.github.io)

Real-World Applications#

Long-context assistants: A coding assistant may process thousands of source-code tokens before generating an answer. GQA reduces the cached key-value state associated with that history, allowing a server to handle longer files or more concurrent users without proportionally increasing GPU memory.

Multimodal image and video assistants: A vision-language model can represent image patches or video frames as long token sequences. GQA reduces attention-cache pressure after these visual tokens enter the language decoder, potentially leaving more memory for additional images, frames, or requests. This differs from a standard Vision Transformer, where attention may process all image patches in parallel without autoregressive KV caching.

Practical Implementation and Trade-Offs#

PyTorch exposes GQA through enable_gqa=True. This documented CUDA example uses eight query heads and two shared key-value heads:

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)

The output retains eight query heads, while keys and values use only two heads. The PyTorch attention backend selector can control which supported kernel performs the operation, while the PyTorch Transformer building-block guide provides broader implementation context. (docs.pytorch.org)

GQA must be chosen when designing or adapting a model; it is not generally an inference-time switch for an existing checkpoint. Developers should verify framework support, head divisibility, numerical behavior, memory use, and task quality. For computer vision deployment, complementary workflows such as the Ultralytics TensorRT integration and Ultralytics benchmark mode help measure whether architectural and runtime optimizations deliver meaningful improvements on the target hardware.