🧪 The Lobotomized Attention Head Bug™ — where one head does all the work while the others stare into the void

[ToddThomson]

🧠 The Lobotomized Attention Head Bug — A Transformer Debugging Rite of Passage

I just finished building the fast prefill → decode inference path in my Mila DNN library.
Everything seemed fine — the model produced coherent text, KV caching worked, and decode mode looked solid.

But… my transformer’s residuals were way off compared to 🤗 HF GPT-2, and the hidden states just felt wrong.
Not exploding, not NaN’ing — just wrong.

After hours of combing through attention math, KV cache, QKV packing, layernorm, and positional encodings…

I found it.

🧪 Root Cause

In the MHA prefill path, my unpermute_output kernel was wrong.
It needed a padded variant (unpermute_output_padded), and instead of writing all attention heads back into the output tensor…

It only wrote back ONE head.
All the other heads?
Nowhere.
Silent.
Forgotten.
Lobotomized.

🤡 Symptoms (that still produced coherent text!)

Hidden states completely misaligned from HF

Residuals with huge swings

Prefill corrupted → Decode still worked (go figure!)

Yet… model still produced coherent sentences
(Transformers are absurdly resilient.)

🔍 Why it still “worked”

Decode path was correct (so per-token incremental attention was fine)

LayerNorm aggressively stabilized everything

MLP + embeddings carried most of the workload

Attention became “Single-Head Attention + Moral Support”

🎉 Lesson

If you ever see:

Prefill mismatch

HF vs your model drifting hard

Residuals acting hyperactive

Yet decode produces intelligible sentences…

Check your unpermute logic.
A single bad stride or head offset can quietly “turn off” most of attention.

[ToddThomson]

🪦 Memorial Style

In loving memory of Attention Heads 1–11.
They attended every forward pass.
They contributed nothing.
They will be missed.