Decoder-Only Transformer: MHA vs GQA vs MQA

Modern decoder-only language models rely on autoregressive attention, where inference repeatedly computes attention over previously generated tokens.

I wanted to understand one specific efficiency trade-off:

How much inference efficiency can we gain by sharing key/value projections across attention heads, and what happens to language-model quality when we do so?

I built a common Transformer architecture and training setup, then implemented MHA, GQA, and MQA as interchangeable attention modules.

The models were evaluated using the same held-out test set and compared across:

• Language-model loss
• Autoregressive generation throughput
• Peak GPU memory

The implementation also includes RoPE, KV caching, custom BPE tokenization, gradient clipping, learning-rate scheduling, checkpointing, and experiment tracking.

The experiments showed a clear efficiency–quality trade-off.

MHA achieved the lowest test loss at 2.376, while MQA reached the highest throughput at 202 tokens/s. GQA provided an intermediate point, reaching 194 tokens/s with a test loss of 2.394.

This made the comparison more useful than looking at model quality alone: reducing KV projections can improve inference efficiency, but the choice of attention mechanism affects both performance and model quality.

What I built

This project is a from-scratch implementation of a small GPT-style decoder-only Transformer in PyTorch.

Rather than relying on a high-level Transformer implementation, I built the core components myself and used the project to understand how architectural and implementation choices affect training and inference.

The final model contains roughly 12–14M parameters and was trained on WikiText-103 using an NVIDIA A40.

Architecture

The model uses:

• Decoder-only Transformer blocks
• Causal self-attention
• Rotary Positional Embeddings (RoPE)
• Feed-forward / MLP layers
• Residual connections
• Layer normalization
• KV caching for autoregressive generation

I also implemented a custom BPE tokenizer and training pipeline around the model.

Attention variants

The main experiment was comparing three attention mechanisms:

Multi-Head Attention (MHA)

Each attention head has its own query, key, and value projections.

Grouped-Query Attention (GQA)

Multiple query heads share key/value heads, reducing the amount of KV state that needs to be maintained during inference.

Multi-Query Attention (MQA)

All query heads share a single key/value head, maximizing KV sharing.

All three variants were integrated into the same model architecture so that the comparison focused on the attention mechanism rather than unrelated architectural differences.

Results

The results show the expected trade-off: sharing K/V projections improves inference efficiency, while MHA retains the lowest held-out loss in this experiment.

GQA was particularly interesting as an intermediate configuration, providing a substantial throughput improvement over MHA while remaining close in test loss.

What I learned

The project ended up being as much about debugging and understanding the implementation as building the architecture itself.

For example, the model initially started training with unexpectedly high loss. Investigating the issue led me to GPT-style residual-aware initialization and a better understanding of how parameter scale interacts with residual connections and network depth.

I also used the project to understand why KV caching matters during autoregressive generation and how MHA, GQA, and MQA change the amount of key/value state maintained during inference.

Technical write-ups

I'll be adding standalone technical articles covering individual parts of the implementation and the problems I encountered while building it.

Each article is intended to be readable independently, while this project page serves as the overview and experiment hub.