Multi-Query Attention › MQA Memory Benefits
Language AI HandbookAll chapters Multi-Query Attention › MQA Memory BenefitsOne section · Sections in this chapter5 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
MQA's main benefit appears during inference. The KV cache reduction is proportional to the number of heads, which can be substantial for modern large models.
Calculating KV Cache Savings
def calculate_kv_cache_mqa(
num_layers: int,
num_heads: int, # Only matters for MHA
head_dim: int,
seq_length: int,
batch_size: int = 1,
bytes_per_element: int = 2,
use_mqa: bool = False,
) -> dict:
"""Calculate KV cache for MHA vs MQA."""
if use_mqa:
# MQA: single K and V regardless of num_heads
num_kv_heads = 1
else:
# MHA: K and V for each head
num_kv_heads = num_heads
kv_cache_bytes = (
2 # K and V
* num_layers
* num_kv_heads
* seq_length
* head_dim
* batch_size
* bytes_per_element
)
return {
"bytes": kv_cache_bytes,
"megabytes": kv_cache_bytes / (1024**2),
"gigabytes": kv_cache_bytes / (1024**3),
"num_kv_heads": num_kv_heads,
}
# Compare for various configurations
configs = [
{"name": "7B-scale", "layers": 32, "heads": 32, "head_dim": 128},
{"name": "13B-scale", "layers": 40, "heads": 40, "head_dim": 128},
{"name": "70B-scale", "layers": 80, "heads": 64, "head_dim": 128},
]
KV Cache Comparison (32K context): Model MHA Cache MQA Cache Reduction Heads -------------------------------------------------------- 7B-scale 16.00 GB 0.500 GB 32x 32 13B-scale 25.00 GB 0.625 GB 40x 40 70B-scale 80.00 GB 1.250 GB 64x 64
The table shows that the cache reduction factor exactly matches the number of attention heads. For the 7B-scale model with 32 heads, MQA reduces the cache from 8 GB to 0.25 GB. The 70B-scale model with 64 heads drops from 40 GB to just 0.625 GB. These reductions determine whether inference is feasible on given hardware.
The cache reduction factor equals the number of attention heads. A 7B model with 32 heads sees its cache shrink from 8 GB to 0.25 GB, while a 70B model with 64 heads drops from 40 GB to under 1 GB. This is the difference between requiring a cluster of GPUs and running inference on a single consumer card.
The scaling behavior becomes even clearer when we plot cache growth as a function of context length.
The linear relationship between context length and cache size explains why long-context models are so memory-hungry. At 128K tokens, the MHA cache alone requires 32 GB, which exceeds the memory of most consumer GPUs before we even load the model weights. MQA brings this down to 1 GB, making long-context inference practical on a single GPU.
Throughput Implications
The memory savings translate directly into throughput improvements. With less memory consumed by the KV cache, you can:
- Serve longer contexts: Fit 128K tokens instead of 8K on the same hardware
- Increase batch sizes: Process more requests in parallel
- Reduce hardware costs: Serve the same load with fewer GPUs
def estimate_max_batch_size(
gpu_memory_gb: float,
model_memory_gb: float,
num_layers: int,
num_heads: int,
head_dim: int,
seq_length: int,
use_mqa: bool,
) -> int:
"""Estimate maximum batch size given GPU memory constraints."""
available_for_cache = gpu_memory_gb - model_memory_gb
cache_per_sequence = calculate_kv_cache_mqa(
num_layers=num_layers,
num_heads=num_heads,
head_dim=head_dim,
seq_length=seq_length,
batch_size=1,
use_mqa=use_mqa,
)["gigabytes"]
max_batch = int(available_for_cache / cache_per_sequence)
return max(1, max_batch)
# Example: 7B model on A100-80GB
model_config = {"num_layers": 32, "num_heads": 32, "head_dim": 128}
model_memory = 14 # FP16 7B model ~14GB
gpu_memory = 80Maximum batch sizes on A100-80GB (7B model, 66GB for cache): Context Length MHA Batch MQA Batch Throughput Gain -------------------------------------------------------- 4,096 tokens 33 1056 32.0x 8,192 tokens 16 528 33.0x 16,384 tokens 8 264 33.0x 32,768 tokens 4 132 33.0x 65,536 tokens 2 66 33.0x
At 64K context, MHA can only serve 1 request at a time, while MQA can batch 32+ requests. This translates to significant throughput improvements in production settings.
The following heatmap visualizes the feasible operating space for both architectures. Each cell shows whether a given batch size and context length combination fits within GPU memory.
The difference is clear. With MHA, the A100-80GB can only handle batch size 1 beyond 32K context, and batch size 4 is already impossible at 16K. With MQA, you can run batch size 64 at 32K context, or batch size 16 even at 128K. This expanded operating space translates directly to higher throughput and better hardware utilization.