Implement single-head scaled dot-product attention without masking. Given query, key, and value matrices Q, K, V, each of shape (N,D), compute the attention output of shape (N,D) in three stages.
First the scaled score matrix:
scores=DQK⊤Then a numerically stable row-wise softmax:
attni,j=∑lexp(scoresi,l−maxkscoresi,k)exp(scoresi,j−maxkscoresi,k)Finally the value aggregation:
output=attnVAll matrices are row-major. The softmax is taken over the last axis (each query row independently).
Input: N = 2, D = 2, Q = [[1, 0], [0, 1]], K = [[1, 0], [0, 1]], V = [[2, 0], [0, 3]]
Output: [[1.3395, 0.9907], [0.6605, 2.0093]]
Each query attends slightly more to its matching key, so the output is a weighted blend of the two value rows.
Input: N = 2, D = 2, Q = [[0, 0], [0, 0]], K = [[1, 2], [3, 4]], V = [[1, 1], [2, 2]]
Output: [[1.5, 1.5], [1.5, 1.5]]
Zero queries give uniform attention, so each output row is the mean of the value rows.
Decompose the problem into three kernels: scores, row-wise softmax, then value aggregation. Allocate an N×N scratch buffer for the scores.
The middle step is exactly the row-wise softmax you have already built; reuse the shared-memory max and sum-exp reduction over each row.
The scale factor is 1.0f / sqrtf((float)D), applied to each score before the softmax.
Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Implement single-head scaled dot-product attention without masking. Given query, key, and value matrices Q, K, V, each of shape (N,D), compute the attention output of shape (N,D) in three stages.
First the scaled score matrix:
scores=DQK⊤Then a numerically stable row-wise softmax:
attni,j=∑lexp(scoresi,l−maxkscoresi,k)exp(scoresi,j−maxkscoresi,k)Finally the value aggregation:
output=attnVAll matrices are row-major. The softmax is taken over the last axis (each query row independently).
Input: N = 2, D = 2, Q = [[1, 0], [0, 1]], K = [[1, 0], [0, 1]], V = [[2, 0], [0, 3]]
Output: [[1.3395, 0.9907], [0.6605, 2.0093]]
Each query attends slightly more to its matching key, so the output is a weighted blend of the two value rows.
Input: N = 2, D = 2, Q = [[0, 0], [0, 0]], K = [[1, 2], [3, 4]], V = [[1, 1], [2, 2]]
Output: [[1.5, 1.5], [1.5, 1.5]]
Zero queries give uniform attention, so each output row is the mean of the value rows.
Decompose the problem into three kernels: scores, row-wise softmax, then value aggregation. Allocate an N×N scratch buffer for the scores.
The middle step is exactly the row-wise softmax you have already built; reuse the shared-memory max and sum-exp reduction over each row.
The scale factor is 1.0f / sqrtf((float)D), applied to each score before the softmax.
Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.