Triton Single-Query Attention Kernel
Problem Statement
Implement scaled dot-product attention for a single query against N keys/values of dimension D (small enough to fit one block). Compute out = softmax(q ยท Kแต / sqrt(D)) @ V, a length-D vector.
Background
This is a compact version of the fused-attention tutorial. Steps: load q (D,) and K (N, D); scores = sum(K * q, axis=1) * scale; stable softmax over the N scores; weighted sum of V (N, D) rows by the attention weights.
Your Task
Implement attention_kernel and run(N=64, D=32) comparing to a torch reference.
How it is tested
Your solution must define a top-level function run(...) that allocates inputs on the GPU, launches your Triton kernel, and returns a boolean from torch.allclose(triton_out, torch_reference, ...). The grader prints run(...); the expected output is True.
Example:
N = 64, D = 32
True
- The input values
N = 64andD = 32are used to allocate inputsq,K, andVon the GPU, withqbeing a vector of lengthDandKandVbeing matrices of sizeN x D. - The attention scores are computed as
scores = sum(K * q, axis=1) * scale, wherescale = 1 / sqrt(D), resulting in a vector of lengthNwith scores representing the similarity betweenqand each row ofK. - A stable softmax function is applied to the attention scores, producing a probability distribution over the
Nkeys: softmax(scores)=โi=1Nโexp(scoresiโ)exp(scores)โ. - The final output is computed as the weighted sum of the rows of
Vusing the softmax probabilities, and the result is compared to a PyTorch reference implementation usingtorch.allclose, yielding the outputTrueif the results match within a certain tolerance.
Constraints:
- Single program (grid = (1,)) handling one query vs N keys
- scale = 1/sqrt(D); stable softmax over the N scores (subtract max)
- out[d] = sum_n weight[n] * V[n, d]
Background Knowledge
The problem involves implementing scaled dot-product attention, a key component in many neural network architectures, particularly in transformer models. This attention mechanism allows the model to focus on specific parts of the input data when generating outputs. The scaled dot-product attention is computed as out=softmax(qโ KT/Dโ)@V, where q is the query vector, K is the key matrix, V is the value matrix, and D is the dimension of the vectors.
The softmax function is used to normalize the attention weights, ensuring they sum up to 1. This is done to prevent extremely large or small values, which can lead to numerical instability. The dot product of the query vector and the key matrix is computed, and then scaled by 1/Dโ to prevent the values from growing too large. The resulting attention weights are then used to compute a weighted sum of the value matrix rows.
In the context of Triton Programming, the problem requires implementing this attention mechanism using a custom kernel, which is a low-level, performance-critical component of the neural network. The kernel must be designed to efficiently compute the attention weights and the weighted sum on a GPU, taking into account the specific memory layout and computational requirements of the hardware.
Algorithm/Approach
The general approach to solving this problem involves the following steps:
- Load the input data, including the query vector, key matrix, and value matrix, onto the GPU.
- Compute the dot product of the query vector and the key matrix, scaling the result by 1/Dโ.
- Apply the softmax function to the resulting attention weights.
- Compute the weighted sum of the value matrix rows using the attention weights.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.