Grouped-Query Attention

DatadogPhone screenHard

You implement grouped-query attention on plain nested lists, with no array library. All tensors are nested lists of floats:

q[b][h][t] is a list of D numbers: query t of query head h, shape [B, Hq, Tq, D].
k[b][g][j] is a list of D numbers: key j of key/value head g, shape [B, Hkv, Tk, D].
v[b][g][j] is a list of Dv numbers: value j of key/value head g, shape [B, Hkv, Tk, Dv].

Hq is divisible by Hkv. Query heads are split into Hkv consecutive groups of Hq / Hkv heads, and every head in a group uses the same key/value head: query head h uses key/value head h // (Hq / Hkv).

Part 1: Attention with shared key/value heads

Implement grouped_query_attention(q, k, v) that computes grouped-query attention and returns the output tensor of shape [B, Hq, Tq, Dv].

For each batch b, query head h and query t, let g = h // (Hq / Hkv) be the key/value head of h:

The score of key j is dot(q[b][h][t], k[b][g][j]) / sqrt(D).
The weights are the softmax of the scores over the Tk keys.
The output row is the weighted sum of the value rows: sum over j of weight_j * v[b][g][j].

There is no mask: every query attends to every key. Round every number in the output with round(x, 6).

Example 1
k
[[[[1,0,0,0],[0,0,0,0]]]]
q
[[[[2,0,0,0]],[[0,0,2,0]]]]
v
[[[[1,2],[3,4]]]]
Output
[[[[1.537883,2.537883]],[[2,3]]]]
Why: Both query heads use the single key/value head. Head 0 scores the two keys at 1 and 0, so its weights are about 0.731059 and 0.268941 and its output is [1.537883, 2.537883]. Head 1 scores both keys at 0, so the weights are 0.5 each and its output is the plain average [2, 3].

Constraints

1 <= B, Hq, Hkv, Tq, Tk, D, Dv <= 16
Hq is divisible by Hkv
-100 <= q[b][h][t][d] <= 100
-100 <= k[b][g][j][d] <= 100
-100 <= v[b][g][j][d] <= 100

Practice this question

Import your own questions

Paste a question you were asked and get a full problem: statement, tests and a checked solution.

Import a question