Skip to content

← Attention & Transformers step 5 of 8

Medium Research

Multi-Query Attention

Multi-query attention (MQA) gives each head its own queries while sharing one set of keys and values across all heads. Different queries can assign different weights to those shared values.

For example, one head might favor the first position while another favors the last. Both read the same value vectors, but their weighted averages can differ.

Your task

Implement multi_query_attention(Q, K, V) for floating-point arrays:

  • Q: shape (seq_len, n_heads, d_k), with a separate query for each head.
  • K: shape (seq_len, d_k), one shared key vector per position.
  • V: shape (seq_len, d_k), one shared value vector per position.

All dimensions are positive. For each head h, compute:

scores = Q[:, h, :] @ K.T / sqrt(d_k)
weights = softmax(scores, axis=-1)
output_h = weights @ V

Scores have shape (seq_len, seq_len). Normalize each row over key positions using a numerically stable softmax.

Return shape (seq_len, n_heads, d_k), with output_h at output[:, h, :]. Stack the head outputs on axis 1; do not flatten the heads into the feature axis. There is no mask, dropout, batch dimension or learned projection in this exercise. Every query can read every position.

Why share keys and values?

For the same sequence length, head width and storage precision, storing one K/V head takes 1 / n_heads of the K/V storage used by one pair per query head. That comparison concerns the keys and values, not total model memory. Sharing them does not by itself guarantee unchanged model quality.

The diagram uses eight query heads and starts with one shared K/V pair. Changing the K/V head count shows how MQA relates to grouped-query and multi-head attention. This exercise specifically implements the one-pair case.

Loading visualization…