MQA Broadcast Outputs

Problem

Run multi-query attention where every query head broadcasts the same K and V matrices. For each query head, compute the standard scaled softmax attention output row by row.

The signature is mqa_broadcast_outputs(Q_heads, K, V). Q_heads is a list of (seq_q, d) query matrices and K, V are (seq_k, d); return a list of (seq_q, d) output matrices.

Hint

All query heads share one set of keys and values, which is exactly what saves cache memory in MQA.

Test cases

case 1

in: [[[[1,0]],[[0,1]]],[[1,0],[0,1]],[[1,2],[3,4]]]

expected: [[[1.6604769013,2.6604769013]],[[2.3395230987,3.3395230987]]]

case 2

in: [[[[1,0]]],[[1,0],[0,1],[1,1]],[[1,0],[0,1],[1,1]]]

expected: [[[0.8022241854,0.5988879073]]]

case 3

in: [[[[1,1]],[[2,0]]],[[1,0]],[[5,6]]]

expected: [[[5,6]],[[5,6]]]

case 4

in: [[[[1,0],[0,1]]],[[1,0],[0,2]],[[1,1],[2,0]]]

expected: [[[1.3302384507,0.6697615493],[1.8044296825,0.1955703175]]]

Starter codePython

import math
def mqa_broadcast_outputs(Q_heads, K, V):
    # Your code here
    pass

4 test cases · run your code against them in the editor