Multi-Query Attention › Practical Implementation Considerations

Michael BrenndoerferUpdated August 4, 20252 min read
Language AI HandbookAll chapters
Part 29 of 67Modern Decoder Models
Chapter 210Multi-Query Attention › Practical Implementation Considerations
Reading section 8 of 11 · Practical Implementation Considerations · 2 minSection 8 of 11 · 2 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 210 of 513

Section 8 of 11 · 2 min read

Multi-Query Attention › Practical Implementation ConsiderationsOne section · Sections in this chapter8 / 11 ↓
Sections in this chapter11
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

Converting MHA to MQA

If you have a pre-trained MHA model and want MQA's efficiency gains, you can convert it through careful weight averaging:

In[24]:
Code
def convert_mha_to_mqa(
    mha_module: nn.Module, num_heads: int, head_dim: int
) -> MultiQueryAttention:
    """
    Convert a pre-trained MHA module to MQA by averaging K and V projections.

    This is a form of knowledge distillation that preserves average behavior
    while reducing parameters.
    """
    hidden_size = mha_module.q_proj.in_features

    # Create MQA module
    mqa = MultiQueryAttention(hidden_size, num_heads, head_dim)

    # Copy Q projection directly
    mqa.q_proj.weight.data = mha_module.q_proj.weight.data.clone()

    # Average K projections across heads
    k_weights = mha_module.k_proj.weight.data
    k_weights = k_weights.view(num_heads, head_dim, hidden_size)
    mqa.k_proj.weight.data = k_weights.mean(dim=0)

    # Average V projections across heads
    v_weights = mha_module.v_proj.weight.data
    v_weights = v_weights.view(num_heads, head_dim, hidden_size)
    mqa.v_proj.weight.data = v_weights.mean(dim=0)

    # Copy output projection
    mqa.o_proj.weight.data = mha_module.o_proj.weight.data.clone()

    return mqa


# Demonstrate the conversion
mha_layer = StandardMultiHeadAttention(
    hidden_size=4096, num_heads=32, head_dim=128
)
# Add output projection for completeness
mha_layer.o_proj = nn.Linear(32 * 128, 4096, bias=False)
Out[25]:
Console
Conversion complete!
Original MHA K projection: torch.Size([4096, 4096])
Converted MQA K projection: torch.Size([128, 4096])
Parameter reduction in K: 32x

The conversion uses averaging, which works reasonably well but may require fine-tuning to recover any lost quality. Research suggests that a small amount of fine-tuning (1-5% of original training) can recover most of the original model's performance.

Uptraining for Better Quality

Rather than training from scratch with MQA, many practitioners start with MHA pre-training and then "uptrain" to MQA or GQA:

  1. Pre-train with MHA: Full expressiveness during the learning phase
  2. Convert to MQA/GQA: Average or group the KV projections
  3. Fine-tune briefly: Recover quality with the new architecture
  4. Deploy efficiently: Enjoy the memory benefits during inference

This approach often yields better results than training MQA from scratch, as the model benefits from the full representational capacity during the critical pre-training phase.