Grouped-Query Attention (GQA)

GQA creates a tunable middle ground between Multi-Head Attention (MHA) and Multi-Query Attention (MQA). It groups multiple query heads (e.g., 8) to share a single key head and a single value head. For example, a model with 32 query heads might be configured with 8 groups, resulting in 8 unique key heads and 8 unique value heads. This grouping factor is a hyperparameter that allows system designers to precisely control the KV cache size versus model expressiveness trade-off.

The primary engineering benefit of GQA is the direct reduction in KV cache memory footprint. The cache size is proportional to the number of unique key and value heads, not the total query heads. This reduction translates to:

• Higher batch sizes for the same GPU memory budget.
• Support for longer context windows.
• Reduced memory bandwidth pressure, alleviating a key bottleneck in the memory-bound decode phase. For a model with h query heads and a group size of g, the KV cache size is reduced by a factor of h/g compared to MHA.

Unlike MQA, which can lead to measurable quality degradation, GQA is designed to preserve model performance. By maintaining multiple independent key/value groups, the model retains more of the representational capacity and diversity of the original MHA. Empirical results, such as those from the Llama 2 and Gemma model families, show that GQA can match the quality of MHA while achieving most of the memory savings of MQA. The grouping allows the model to learn specialized key/value representations for different semantic or syntactic functions.

GQA optimizes both compute and memory access patterns during the decode phase.

• Compute: The attention operation requires fewer matrix multiplications involving the key and value projections, leading to faster computation per decoding step.
• I/O: The smaller KV cache means less data must be read from GPU memory for each attention calculation, which is critical in memory-bandwidth-bound scenarios. This efficiency directly improves token generation latency and increases overall inference throughput.

GQA can be implemented via a simple reshaping operation after the initial linear projections. It is compatible with other core inference optimizations:

• PagedAttention (vLLM): The reduced per-head cache size allows for more efficient block management.
• FlashAttention: Optimized kernels can be adapted for the grouped computation pattern.
• Rotary Positional Embedding (RoPE): Applied normally to the queries and the reduced set of keys.
• Quantization: KV cache quantization techniques are even more effective due to the smaller total cache size. Frameworks like TensorRT-LLM provide native support for compiling models with GQA.

The group size g is the central lever for system designers. It defines a continuous spectrum:

• g = 1 is Multi-Query Attention (MQA): Maximum memory savings, potential quality loss.
• g = h (number of query heads) is Multi-Head Attention (MHA): No memory savings, original quality.
• 1 < g < h is GQA: Configurable balance. Choosing g involves benchmarking target workloads to find the optimal point where further cache reduction provides diminishing returns on latency/throughput versus any measurable drop in task accuracy. This makes GQA a key tool for inference cost optimization.

The standard attention mechanism in transformers where each attention head maintains its own independent set of query (Q), key (K), and value (V) projection matrices. This provides high model capacity and quality but results in a large KV cache size, as the cache must store K and V vectors for every head and every token.

• KV Cache Size: Scales linearly with the number of attention heads.
• Baseline for Comparison: GQA and MQA are evaluated against MHA's performance benchmark.

An attention variant where all query heads share a single key head and a single value head. This drastically reduces the size of the KV cache, as only one set of K and V vectors is stored per layer, regardless of the number of query heads.

• Primary Benefit: Maximum reduction in KV cache memory and memory bandwidth.
• Trade-off: Can lead to a measurable drop in model quality compared to MHA, as the model's capacity to represent different aspects of the input is reduced.
• Use Case: Found in models like Falcon, where inference efficiency is paramount.

The memory buffer that stores pre-computed key and value tensors from previous tokens during the autoregressive decode phase of a transformer. Instead of recomputing these tensors for every new token, the model reads them from this cache, which is the primary source of inference latency and memory consumption for long sequences.

• Core Bottleneck: Reading the KV cache is often a memory-bound operation.
• Optimization Target: Techniques like GQA, MQA, and quantization directly target reducing the size and bandwidth requirements of this cache.

A memory management algorithm that organizes the KV cache into non-contiguous, fixed-size blocks or "pages." Popularized by the vLLM inference engine, it treats GPU memory like an operating system's virtual memory.

• Solves Fragmentation: Enables efficient sharing of physical memory between different requests in a continuous batch.
• Enables GQA/MQA: Efficiently manages the variable-sized KV caches produced by these attention mechanisms when requests have different context lengths.
• Key Innovation: Allows near-zero waste in KV cache memory allocation.

An inference scheduling technique that dynamically groups incoming requests into a single, continuously running batch. New requests can join the batch, and finished sequences can exit, independently.

• Maximizes GPU Utilization: Eliminates idle time caused by waiting for a static batch to finish.
• Synergy with GQA: The reduced per-request KV cache footprint of GQA allows continuous batching to pack more concurrent requests into the available GPU memory, driving higher overall throughput.

A memory compression technique that stores the key and value tensors in a lower numerical precision format, such as INT8 or FP8, instead of the standard FP16/BF16.

• Direct Memory Reduction: Can cut the KV cache memory footprint by 50% or more.
• Complementary to GQA: While GQA reduces the number of K/V vectors stored, quantization reduces the size of each vector. They are often used together for maximum memory savings.
• Trade-off: Introduces potential quantization error, which must be managed to preserve model quality.