Grouped-Query Attention: Which Group Is a Head In?
Beginner
The problem, from first principles
Grouped-query attention (Ainslie et al., 2023) sits between multi-head and multi-query attention. Query heads are split into groups, and each group shares one key-value head. The grouping is a simple integer division.
From theory to code
Implement group_index(head, n_heads, n_kv), returning the key-value head used by a query head.
Constraints
- Query heads are assigned to groups in contiguous blocks.
Hints
Hint 1
Each group has n_heads / n_kv query heads, so divide the head index by that group size.
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.