Grouped-Query Attention

Asked at Datadog· Screen

Reported by the problem's creator, not verified

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).

Example 1
Input: q = [[[[2,0,0,0]],[[0,0,2,0]]]], k = [[[[1,0,0,0],[0,0,0,0]]]], v = [[[[1,2],[3,4]]]]
Output: [[[[1.537883,2.537883]],[[2,3]]]]
Explanation: 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].

Import your own questions

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

Import a question

Grouped-Query Attention

Examples

Example 1

Code