Multi-Query Attention: Shared Keys and Values
Advanced
The problem, from first principles
In multi-query attention every query head attends to the same keys and values. Each head keeps its own queries, so the model still learns different attention patterns, but only one set of keys and values is stored and read.
From theory to code
Implement multi_query_attention(q, k, v), computing attention for all heads against the shared keys and values.
Constraints
kandvhave no head axis.
Hints
Hint 1
Broadcast the shared keys across heads by using matrix products with the head axis first, then softmax and weight the shared values.
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.