Grouped-Query Attention: Expanding Key-Value Heads
Intermediate
The problem, from first principles
In grouped-query attention, each key-value head serves a group of query heads. Before the attention product, the key and value heads are repeated so that every query head finds its group's copy.
From theory to code
Implement expand_kv(kv, n_heads), repeating each key-value head across its group.
Constraints
n_headsis a multiple of the number of key-value heads.
Hints
Hint 1
Repeat along the head axis by the group size, keeping groups in order.
Keywords
research-paperssystemsattentioninference
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.