Build Causal Multi-Query Attention (MQA)

Hard · Solved in PyTorch · Machine learning coding practice

Problem

Implement the forward pass for a Multi-Query Attention (MQA) layer. This optimization drastically speeds up autoregressive decoding in LLMs by reducing the memory footprint of the KV cache.

Unlike standard Multi-Head Attention (MHA) where Query, Key, and Value all have H heads, MQA utilizes H Query heads but only 1 shared Key head and 1 shared Value head.

Topics: transformers, attention, pretrain

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Build Causal Multi-Query Attention (MQA) in the browser and run it against the tests.

Related problems

• Build Causal Grouped-Query Attention (GQA) Hard
• Build Causal Multi-Head Attention for nanoGPT Hard
• Calculate Self-Attention Easy
• MoE: Dispatch and Weighted Combine Medium
• MoE: Router Load-Balancing Loss Medium
• MoE: Top-k Expert Routing Easy