Implement Grouped Query Attention from Scratch in PyTorch

Mediummodern-architecturesauto-graded

Build GQA where multiple query heads share key-value heads, reducing KV cache memory while preserving quality.

Solve it

Open in ColabQuestion notebookSolution

Check your answer

The grader verifies properties of your implementation, so a correct solution written differently from ours still passes.

pip install torchleet

from torchleet import check
check("grouped-query-attention", grouped_query_attention)