Multi-Head Attention Forward Pass
Meta● HardFULLTIMEPHONE SCREENONSITE INTERVIEWProblem statement
Implement a projection-free multi-head scaled dot-product attention forward pass. query, key, and value are row-major integer matrices with the same feature width d; key and value have the same number of rows. heads divides d.
Split each row into contiguous head slices. For each query row and head, compute dot products against every key slice, divide by sqrt(d / heads), apply a numerically stable softmax, and use the weights to combine the matching value slices. Concatenate the head outputs.
Return one comma-separated string per query row with every coordinate formatted to exactly six decimal places. There are no learned projections or masks.
Interview discussion (not judged by the function)
After implementing the forward pass, analyze its time and extra-space complexity in terms of query rows q, key/value rows k, feature width d, and head count h. Explain conceptually how FlashAttention reduces GPU memory traffic while computing exact softmax attention, and how linear attention changes the sequence-length cost by using a factorized feature-map formulation. No implementation of either alternative is required.
Function
multiHeadAttention(query: int[][], key: int[][], value: int[][], heads: int) → String[]Examples
Example 1
query = [[1]]key = [[0],[1]]value = [[10],[20]]heads = 1return = ["17.310586"]Softmax over scores 0 and 1 weights the second value more heavily.
Example 2
query = [[1,0]]key = [[1,0],[0,1]]value = [[4,8],[6,2]]heads = 2return = ["4.537883,5.000000"]Each one-dimensional head computes its own attention distribution.
Constraints
1 <= query rows, key rows, d <= 40.- All rows have width
d;1 <= heads <= dand dividesd. - Entries are between -100 and 100.