Why is implementing multi-head attention with a for loop slow?

Sreedath Panat times a loop-based multi-head attention wrapper in Google Colab: one forward pass took 0.003 seconds with 2 heads, about an order of magnitude longer with 400 heads, and 0.63 seconds with 4,000 heads. The reason he gives is that every head adds its own separate matrix multiplication, and a GPU handles many small separate multiplications far less efficiently than one large one, since each has to be shipped, loaded and returned. He adds that the latency hits inference as well as training.

• The loop version builds one causal attention object per head in an nn.ModuleList and joins outputs with torch.cat on the last dimension.
• Panat admits the 400 and 4,000 head runs break the usual d over n sizing and serve only to show the slowdown.

Source: our own transcription of the episode. The full conversation is on The attention family (2026-10-03). Read the whole episode summary →

▸More from this episode