Attention Topologies MHA GQA MQA and MLA

Every token generated during autoregressive decoding requires accessing the Key and Value states of all previously generated tokens in the sequence. Storing these historical states in GPU memory prevents the model from recomputing them at every step. This storage mechanism is the KV cache. As context windows expand from 4K to 128K tokens and batch sizes increase to saturate GPU memory bandwidth, the physical size of the KV cache rapidly becomes the primary bottleneck for inference scaling. Modifying the internal routing of attention layers provides a direct architectural solution to this memory pressure.

The exact physical memory footprint of the KV cache dictates how many concurrent requests a specific hardware configuration can support. You calculate this footprint using the following equation:

$Memory_{KV} = 2 \times num\_layers \times num\_kv\_heads \times head\_dim \times seq\_length \times bytes\_per\_element \times batch\_size$

The multiplier of 2 accounts for storing both the Key and the Value tensors. To understand how architectural choices alter this footprint, we evaluate how different attention structures modify the num_kv_heads variable and change the memory demands on enterprise accelerators like the A100 80GB or consumer hardware like the RTX 4090 24GB.

Multi-Head Attention Topologies

Multi-Head Attention (MHA) is the standard baseline established by the original transformer architecture. In MHA, the number of Key and Value heads is exactly equal to the number of Query heads. Every individual Query head processes its own dedicated, independent Key and Value representations.

While this provides maximum expressivity, it scales poorly in production inference. We can quantify this by calculating the KV cache size for a model utilizing MHA, such as meta-llama/Llama-2-7b-chat-hf. This model has 32 layers, 32 KV heads, a head dimension of 128, and typically operates in FP16 precision (2 bytes per element).

For a single user sequence of 4,096 tokens: $Memory_{KV} = 2 \times 32 \times 32 \times 128 \times 4096 \times 2 \times 1 = 2,147,483,648 \text{ bytes}$

A single 4K context sequence consumes approximately 2.15 GB of VRAM. If you run a batch size of 16 to maximize GPU compute utilization, the KV cache alone requires 34.3 GB. The model weights in FP16 consume about 14 GB. The combined total of 48.3 GB immediately triggers an Out-Of-Memory (OOM) error on a 24GB RTX 3090 or RTX 4090, forcing infrastructure operators to shard the model across multiple GPUs or implement heavy quantization.

Grouped-Query and Multi-Query Attention

To alleviate the severe memory constraints of MHA, model designers introduced Multi-Query Attention (MQA) and Grouped-Query Attention (GQA). Both topologies reduce the num_kv_heads variable in the cache equation.

Multi-Query Attention reduces the number of KV heads to exactly 1. All Query heads share the same single Key and Value projection. This reduces the KV cache size by a factor of $N$ (where $N$ is the number of query heads), but often results in measurable degradation in model reasoning quality.

Grouped-Query Attention acts as a balanced middle ground. It divides the Query heads into discrete groups, and each group shares a single Key and Value head. Modern architectures heavily favor GQA. For example, meta-llama/Meta-Llama-3.1-8B employs GQA with 32 Query heads and 8 KV heads. This means 4 Query heads share a single KV head.

We can recalculate the exact memory footprint for Llama 3.1 8B using the same 4,096 token context length: $Memory_{KV} = 2 \times 32 \times 8 \times 128 \times 4096 \times 2 \times 1 = 536,870,912 \text{ bytes}$

The cache per sequence drops from 2.15 GB to roughly 0.54 GB. For a batch size of 16, the total KV cache requires only 8.6 GB. Added to the ~16 GB model weights, the total memory footprint sits at 24.6 GB. With minor 8-bit KV cache quantization or context length trimming, this production workload fits comfortably on a single 24GB GPU, entirely due to the GQA topology.

Routing patterns comparing the 1:1 allocation in MHA, the intermediate sharing in GQA, and the fully centralized KV state in MQA.

Multi-Head Latent Attention

Multi-Head Latent Attention (MLA) represents a newer optimization designed for extremely long context windows and large-scale serving. Models like DeepSeek V2 and V3 utilize this approach to compress the memory footprint more than what GQA achieves.

Instead of caching explicitly projected Key and Value tensors for every token, MLA projects the input representations into a low-dimensional joint latent vector, denoted as $c_t$, with dimension $d_c$. The model only stores this compressed latent vector in the KV cache. During the autoregressive decode phase, the model dynamically decompresses $c_t$ back into the separate Key and Value representations using frozen up-projection matrices.

If the latent dimension $d_c$ is 512, storing a single vector of 512 elements replaces storing multiple KV heads each with a dimension of 128. For a 32-head model, MHA would store $2 \times 32 \times 128 = 8,192$ elements per token. MLA stores 512 elements. This yields a massive reduction in memory overhead.

On an enterprise cluster using A100 80GB hardware, generating a 128,000 token sequence with standard GQA often requires distributing the KV cache across multiple nodes using Tensor Parallelism or Ring Attention. By applying MLA, infrastructure operators can fit the entire 128K context cache directly into the VRAM of a single accelerator, severely reducing interconnect latency across NVLink boundaries.

Inspecting Model Topologies in Code

When sizing infrastructure, you can verify the attention topology of any open-source model programmatically. The Hugging Face transformers library exposes the num_attention_heads (Queries) and num_key_value_heads directly in the configuration object.

Inspecting the model configuration directly provides the exact ratio of Query heads to KV heads, defining GQA hardware requirements before loading weights into VRAM:

from transformers import AutoConfig

# Load the configuration without downloading the model weights
model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
config = AutoConfig.from_pretrained(model_id)

q_heads = config.num_attention_heads
kv_heads = config.num_key_value_heads
head_dim = config.hidden_size // q_heads
layers = config.num_hidden_layers

print(f"Model: {model_id}")
print(f"Query Heads: {q_heads}")
print(f"KV Heads: {kv_heads}")
print(f"Head Dimension: {head_dim}")
print(f"GQA Ratio: {q_heads // kv_heads} Query heads per KV head")

# Calculate the per-token KV cache size in bytes (assuming FP16)
bytes_per_element = 2
bytes_per_token = 2 * layers * kv_heads * head_dim * bytes_per_element

print(f"KV Cache per token: {bytes_per_token} bytes")

Running this code reveals that Mixtral 8x7B utilizes 32 Query heads and 8 KV heads, a 4:1 GQA ratio identical to Llama 3.1 8B. It consumes precisely 131,072 bytes (128 KB) per token in the sequence length. You multiply this base per-token metric by your maximum expected batch size and target sequence length to provision your physical VRAM accurately.