Triton Single-Query Attention Kernel

hardPROAdvanced KernelsTriton Programming

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

Input

N = 64, D = 32

Output

True

Reasoning

The input values N = 64 and D = 32 are used to allocate inputs q, K, and V on the GPU, with q being a vector of length D and K and V being matrices of size N x D.
The attention scores are computed as scores = sum(K * q, axis=1) * scale, where scale = 1 / sqrt(D), resulting in a vector of length N with scores representing the similarity between q and each row of K.
A stable softmax function is applied to the attention scores, producing a probability distribution over the N keys: $softmax(scores) = \frac{exp(scores)}{\sum_{i=1}^{N} exp(scores_i)}$ .
The final output is computed as the weighted sum of the rows of V using the softmax probabilities, and the result is compared to a PyTorch reference implementation using torch.allclose, yielding the output True if 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]

Learn this in

Advanced Kernels & OptimizationGPU Basics with Triton

More from Triton Programming

Triton Row Cumulative Sumhard Triton Stable LogSumExp Kernelhard Triton Seeded Fused Dropouthard Triton Autotuned Vector Addmedium Triton Matrix-Vector Product (GEMV)medium

Concept primerPREVIEW

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 = \text{softmax}(q \cdot K^T / \sqrt{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/\sqrt{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/\sqrt{D}$ .
Apply the softmax function to the resulting attention weights.
Compute the weighted sum of the value matrix rows using the attention weights.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.