Multi-head attention with a KV cache and grouped-query attention

Implement the attention layer of a decoder-only transformer from scratch in three steps that build on each other:

1. Standard multi-head attention with a causal mask.
2. A key-value cache so generation can add one token at a time without recomputing earlier keys and values.
3. The forward pass of grouped-query attention, where several query heads share one key-value head, using the cache.

Keep the tensor shapes explicit and say how much memory the cache saves in step 3.