Multi-Query Attention › Implementing Multi-Query Attention
Language AI HandbookAll chapters Multi-Query Attention › Implementing Multi-Query AttentionOne section · Sections in this chapter4 / 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
Now that we understand the mathematical formulation, let's translate it into working code. The implementation will reveal exactly how the parameter sharing manifests in practice and why it leads to such significant memory savings.
We'll build MQA from scratch, examining each component to understand how it differs from standard multi-head attention. By the end, you'll see that the core change is just a matter of adjusting projection dimensions.
The Core MQA Module
The main difference between MHA and MQA lies in the output dimensions of the projection layers. In MHA, the key and value projections output num_heads * head_dim dimensions (one set of K/V vectors per head). In MQA, they output just head_dim (a single set of K/V vectors shared across all heads).
Let's implement this step by step, with comments highlighting the critical differences.
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiQueryAttention(nn.Module):
"""
Multi-Query Attention: all heads share a single K and V projection.
This reduces KV cache size by a factor of num_heads while maintaining
expressive power through multiple query heads.
"""
def __init__(
self,
hidden_size: int,
num_heads: int,
head_dim: int | None = None,
dropout: float = 0.0,
):
super().__init__()
self.hidden_size = hidden_size
self.num_heads = num_heads
self.head_dim = head_dim or (hidden_size // num_heads)
self.dropout = dropout
# Each head has its own query projection
# Total query params: hidden_size -> num_heads * head_dim
self.q_proj = nn.Linear(
hidden_size, num_heads * self.head_dim, bias=False
)
# Single shared key projection: hidden_size -> head_dim
self.k_proj = nn.Linear(hidden_size, self.head_dim, bias=False)
# Single shared value projection: hidden_size -> head_dim
self.v_proj = nn.Linear(hidden_size, self.head_dim, bias=False)
# Output projection
self.o_proj = nn.Linear(
num_heads * self.head_dim, hidden_size, bias=False
)
self.scale = self.head_dim**-0.5
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
use_cache: bool = False,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
batch_size, seq_len, _ = hidden_states.shape
# Project queries: (batch, seq, num_heads * head_dim)
q = self.q_proj(hidden_states)
# Project keys and values: (batch, seq, head_dim)
k = self.k_proj(hidden_states)
v = self.v_proj(hidden_states)
# Reshape queries for multi-head: (batch, num_heads, seq, head_dim)
q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
q = q.transpose(1, 2)
# Keys and values stay as single head: (batch, 1, seq, head_dim)
k = k.unsqueeze(1)
v = v.unsqueeze(1)
# Handle KV cache
if past_key_value is not None:
past_k, past_v = past_key_value
k = torch.cat([past_k, k], dim=2)
v = torch.cat([past_v, v], dim=2)
present_key_value = (k, v) if use_cache else None
# Compute attention scores: (batch, num_heads, seq_q, seq_kv)
# Broadcasting: Q is (batch, heads, seq, dim), K is (batch, 1, seq, dim)
# The single K is broadcast across all heads
attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scale
# Apply causal mask if needed
if attention_mask is not None:
attn_weights = attn_weights + attention_mask
# Softmax and dropout
attn_weights = F.softmax(attn_weights, dim=-1)
attn_weights = F.dropout(
attn_weights, p=self.dropout, training=self.training
)
# Apply attention to values
# attn_weights: (batch, heads, seq_q, seq_kv)
# v: (batch, 1, seq_kv, dim) -> broadcasts to all heads
attn_output = torch.matmul(attn_weights, v)
# Reshape and project output
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.view(batch_size, seq_len, -1)
output = self.o_proj(attn_output)
return output, present_key_valueSeveral implementation details deserve attention:
Projection dimensions: The query projection outputs num_heads * head_dim dimensions because each head needs its own query vectors. But the key and value projections output only head_dim, producing a single set of K/V vectors. This is the source of our memory savings.
Broadcasting: When we compute torch.matmul(q, k.transpose(-2, -1)), PyTorch's broadcasting rules handle the dimension mismatch. The query tensor has shape (batch, num_heads, seq, head_dim) while the key tensor has shape (batch, 1, seq, head_dim). Broadcasting automatically replicates the single key across all heads, computing the correct attention scores without physically copying the data.
KV cache handling: The cache stores only the shared K and V tensors, not per-head copies. When we concatenate past keys with new keys, we're working with tensors of shape (batch, 1, seq, head_dim) rather than (batch, num_heads, seq, head_dim). This is where the memory reduction materializes during autoregressive generation.
Comparing Parameter Counts
To make the savings concrete, let's count parameters in MHA versus MQA for a realistic configuration.
def count_parameters(module: nn.Module) -> int:
"""Count trainable parameters in a module."""
return sum(p.numel() for p in module.parameters() if p.requires_grad)
class StandardMultiHeadAttention(nn.Module):
"""Standard MHA for comparison."""
def __init__(
self, hidden_size: int, num_heads: int, head_dim: int | None = None
):
super().__init__()
self.num_heads = num_heads
self.head_dim = head_dim or (hidden_size // num_heads)
total_dim = num_heads * self.head_dim
# All three projections are full-sized
self.q_proj = nn.Linear(hidden_size, total_dim, bias=False)
self.k_proj = nn.Linear(hidden_size, total_dim, bias=False)
self.v_proj = nn.Linear(hidden_size, total_dim, bias=False)
self.o_proj = nn.Linear(total_dim, hidden_size, bias=False)
# Create instances with typical LLaMA-like config
hidden_size = 4096
num_heads = 32
head_dim = 128
mha = StandardMultiHeadAttention(hidden_size, num_heads, head_dim)
mqa = MultiQueryAttention(hidden_size, num_heads, head_dim)Configuration: hidden_size=4096, num_heads=32, head_dim=128 Parameter counts: Standard MHA: 67,108,864 parameters Multi-Query: 34,603,008 parameters Reduction: 48.4% Projection breakdown (MHA): Q: 4096 × 4096 = 16,777,216 K: 4096 × 4096 = 16,777,216 V: 4096 × 4096 = 16,777,216 Projection breakdown (MQA): Q: 4096 × 4096 = 16,777,216 K: 4096 × 128 = 524,288 V: 4096 × 128 = 524,288
The numbers reveal the trade-off clearly. MQA reduces the K and V projection parameters by exactly num_heads (32x in this case). However, the total attention layer parameter reduction is more modest, around 50%, because the query and output projections remain unchanged at full size.
This distinction matters: the parameter savings affect model size and training memory, but the real win comes during inference. The KV cache stores the outputs of the K and V projections, not the projection matrices themselves. Since each forward pass produces only one K vector and one V vector per token (instead of 32), the cache reduction matches the full 32x factor. This is why MQA changes the economics of inference even though its impact on total parameter count is more modest.