Grouped-Query Attention

Task

Implement class GroupedQueryAttention(w_query: list[list[float]], w_key: list[list[float]], w_value: list[list[float]], w_output: list[list[float]], num_query_heads: int, num_kv_heads: int). Build a reusable grouped-query attention layer whose forward method projects query heads separately, shares each projected key-value head across one contiguous query-head group, computes stable masked attention, concatenates the query-head outputs, and applies the output projection.

Requirements

• Store the four projection matrices and both head counts on the instance; copy mutable inputs so constructor arguments are not mutated through the layer.
• In the PyTorch profile, inherit nn.Module, call super().__init__(), and register every projection copy as an nn.Parameter so optimizers and state_dict() can discover it.
• query, key, and value have model width d_model; w_query and w_output have shape d_model by d_model.
• d_model must divide evenly by num_query_heads; each head has width d_head = d_model / num_query_heads.
• num_query_heads must be divisible by num_kv_heads; w_key and w_value have shape d_model by num_kv_heads * d_head.
• Project query, key, and value before splitting their last dimensions into contiguous heads.
• Map query head h to key-value head floor(h / (num_query_heads / num_kv_heads)), so every contiguous query-head group shares one key-value head.
• Within each query head, divide query-key dot products by sqrt(d_head).
• A True mask entry is valid and a False entry is blocked; the same query-key mask applies to every query head.
• Use row-wise stable softmax over valid keys and return zeros for every fully masked query row.
• Concatenate query-head outputs in their original order and multiply by w_output.
• Preserve an empty query matrix by returning an empty list.

Example

attention = GroupedQueryAttention(
    identity, kv_projection, kv_projection, identity, 4, 2,
)
attention.forward(query, key, value, [[True, True]])
# [[7.31058579, 8.80797078, 95.25741268, 98.20137900]]