Multi-Query Attention (MQA)

The defining architectural feature of MQA. Unlike standard Multi-Head Attention (MHA), where each of the N query heads has its own dedicated key and value head, MQA uses a single, shared key head and a single, shared value head for all query heads.

• Mechanism: During the attention computation, the N distinct query vectors all attend to the same, single set of key and value vectors.
• Impact: This reduces the number of key and value tensors that must be stored in the KV cache from 2 * N per layer to just 2 per layer (one key, one value).

MQA's primary benefit is a linear reduction in the memory footprint of the KV cache, which is the dominant memory consumer during autoregressive decoding.

• Quantitative Reduction: For a model with h attention heads, MQA reduces the size of the KV cache by a factor of approximately h. For a 32-head model, this is a ~32x reduction in the number of cached key/value vectors.
• Consequence: This allows for either longer context windows within the same GPU memory budget or higher batch sizes to increase throughput, directly addressing inference cost and latency.

By caching fewer tensors, MQA significantly decreases the amount of data that must be read from GPU memory during each decoding step, moving the computation from a memory-bound regime towards a compute-bound one.

• The Bottleneck: In standard MHA, loading the massive KV cache for attention is often the performance limiter.
• MQA's Effect: With a smaller cache, the memory I/O volume per token is drastically lower. This reduces latency and increases the theoretical maximum tokens-per-second throughput, especially on memory-bandwidth-constrained hardware.

The memory efficiency of MQA comes with a trade-off. Sharing keys/values across all queries can reduce the model's representational capacity and expressiveness.

• The Risk: The shared key/value projection may act as an information bottleneck, potentially leading to a drop in output quality or factual accuracy compared to the full MHA baseline, particularly on complex reasoning tasks.
• Empirical Observation: The degradation is often model and task-dependent. Some large models (e.g., certain versions of Falcon) use MQA effectively, while others require careful tuning or prefer the hybrid Grouped-Query Attention (GQA) approach.

MQA simplifies the attention computation graph and parameter count, leading to cleaner and faster kernel implementations.

• Fewer Parameters: The model has fewer linear projection matrices for keys and values, slightly reducing the total parameter count.
• Kernel Optimization: The simpler data access pattern (many queries to one key/value set) can be optimized into highly efficient GPU kernels. This contributes to lower latency beyond just the cache savings.
• Decoding Step: The autoregressive step becomes: Attn = softmax(Q_all @ K_shared.T) @ V_shared, where Q_all is the concatenated query tensor.

Grouped-Query Attention (GQA) is a generalization and intermediate point between MHA and MQA, created to mitigate MQA's quality trade-off.

• How it Works: Instead of N query heads sharing one key/value head (MQA) or each having its own (MHA), GQA groups queries into G groups. Each group shares a single key head and value head.
• The Spectrum:
  • G = N -> Multi-Head Attention (largest cache, highest quality).
  • 1 < G < N -> Grouped-Query Attention (balanced trade-off).
  • G = 1 -> Multi-Query Attention (smallest cache, highest efficiency).
• Practical Use: GQA allows practitioners to tune the KV cache size versus model quality based on their specific deployment constraints.

Grouped-Query Attention is a hybrid attention mechanism that serves as a middle ground between Multi-Head Attention (MHA) and Multi-Query Attention (MQA). Instead of a single key and value head for all queries (MQA) or independent heads for each query (MHA), GQA groups multiple query heads to share a single key head and value head.

• Design Trade-off: It offers a tunable parameter—the group size—allowing practitioners to balance the KV cache memory reduction of MQA against the model quality and capacity of full MHA.
• Adoption: Models like Llama 2 and Gemma use GQA, as it often recovers most of the quality of MHA while providing significant memory savings over it.
• Implementation: The grouped queries compute attention against a shared, smaller set of key and value projections, reducing the cache size proportionally to the group size.

A KV cache, or key-value cache, is the fundamental memory structure optimized by MQA. During the autoregressive decode phase of a transformer, the model computes key (K) and value (V) tensors for each input token.

• Purpose: Instead of recomputing these tensors for every new token, they are stored in the KV cache. For the next token, the model only computes the new query (Q) and performs attention against the cached K and V from all previous tokens.
• Memory Footprint: The size of the KV cache scales with batch size, sequence length, and the number of attention heads. It is often the primary bottleneck for long-context or high-throughput inference.
• Impact of MQA: MQA directly reduces this footprint by sharing a single K and V head across all query heads, drastically shrinking the cache size compared to standard multi-head attention.

PagedAttention is a memory management algorithm, central to the vLLM inference engine, that optimizes how the KV cache is allocated in physical memory. It treats the cache analogously to virtual memory in an operating system.

• Mechanism: The KV cache for each sequence is divided into fixed-size blocks or pages. These pages can be non-contiguous in physical GPU memory.
• Key Benefit: This eliminates internal fragmentation, a major source of memory waste in continuous batching where sequences have different lengths. It allows for near-optimal memory utilization.
• Synergy with MQA: While MQA reduces the size of the cache per head, PagedAttention optimizes its allocation and sharing across requests. They are complementary techniques for maximizing throughput.

Continuous batching (or iterative batching) is a scheduling optimization that dynamically groups incoming inference requests to maximize GPU utilization, creating a direct interaction with KV cache management.

• Dynamic Scheduling: Unlike static batching, new requests can join the batch as soon as resources are available, and finished sequences can exit independently.
• Cache Implications: This results in a heterogeneous mix of sequence lengths running concurrently. Efficiently managing the disparate KV caches for these sequences is critical; techniques like PagedAttention were invented to solve this challenge.
• System-Level Goal: The combined use of MQA (smaller caches) and advanced memory management (e.g., PagedAttention) enables continuous batching to achieve very high throughput without running out of memory.

KV cache quantization is a compression technique that reduces the memory footprint of the cache by storing the key and value tensors in a lower numerical precision format.

• Precision Trade-off: Common techniques involve storing cache tensors in FP8, INT8, or even INT4 precision instead of the standard FP16 or BF16 used during computation.
• Objective: The goal is to reduce memory bandwidth pressure and increase effective cache capacity, allowing for longer contexts or larger batches. This is a complementary approach to MQA.
• Challenge: Aggressive quantization can degrade model quality if not applied carefully. Quantization-aware scaling factors are often used to minimize the error introduced by the lower precision.

A memory-bound regime describes the state where an inference system's performance is limited by the speed of memory accesses rather than by raw computational power (compute-bound).

• Primary Culprit: The KV cache is a major driver of memory-bound behavior. During each decoding step, the entire cached context must be read from GPU memory to compute attention scores.
• Impact of MQA: By drastically reducing the size of the KV cache, MQA directly alleviates this bottleneck. It decreases the total volume of data that must be fetched per token, improving decode latency and reducing memory bandwidth saturation.
• System Design: Understanding this regime is crucial for CTOs and systems architects. Optimizations like MQA, quantization, and kernel fusion (e.g., FlashAttention) are all targeted at moving the workload away from being memory-bound.