Build Causal Multi-Head Attention for nanoGPT
Hard · Solved in PyTorch · Machine learning coding practice
Problem
Implement the forward pass for a GPT-style Multi-Head Attention (MHA) layer. This module is the "engine" of a Decoder-only Transformer.
In high-performance implementations like nanoGPT, we don't create three separate linear layers for Q, K, and V. Instead, we use one large fused projection for efficiency and then split the resulting tensor.
Topics: transformers, llm-internals, nanogpt, pretrain
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Build Causal Multi-Head Attention for nanoGPT in the browser and run it against the tests.