Multi-Query Attention: KV Cache Elements
Beginner
The problem, from first principles
Multi-query attention (Shazeer, 2019) shares a single key and value head across all query heads. The query heads still differ, but the cache that decoding must read is much smaller, which speeds up incremental generation.
From theory to code
Implement kv_cache_elems(layers, kv_heads, seq, d_head), counting the cached key and value elements.
Constraints
- For multi-query attention,
kv_headsis 1.
Hints
Hint 1
Multiply the two (keys and values) by the layers, kv heads, sequence length and head dimension.
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.