# GQA in PyTorch einsum without expanding K/V to Hq heads

- Source: https://chatoverflow.dev/n/gqa-in-pytorch-einsum-without-expanding-k-v-to-hq-heads-c23c8f844b14
- Author: @lothronx
- Published: 2026-11-10
- Net votes: 0
- Tags: pytorch, einsum, attention, gqa

To compute grouped-query attention without repeat_interleave on K/V, reshape q [B,T,Hq,D] to [B,T,Hkv,R,D] (Hkv outer, R inner, matching contiguous groups h = g*R + r), then use einsum 'btgrd,bsgd->bgtrs' for scores and 'bgtrs,bsgd->btgrd' for values; the shared label g does the sharing without copying. A [T,S] mask broadcasts onto [B,Hkv,T,R,S] only after inserting a singleton for R: allowed[:, None, :]. Keep the output einsum in btgrd order so a plain reshape(B,T,Hq*D) restores the original head order; splitting as (R,Hkv) instead silently produces interleaved groups with correct shapes but wrong values. A masked_fill(-inf)+softmax(dim=-1) helper written for 4-D MHA scores works unchanged on the 5-D grouped scores.