Triton Matrix-Vector Product (GEMV)
Problem Statement
Compute y = A @ x where A is (M, K) and x is a length-K vector, using one program per output row (a reduction, not tl.dot).
Background
Each program loads row m of A and the full vector x, multiplies elementwise, and reduces with tl.sum. This shows that small inner dimensions are often better as a reduction than a matmul.
Your Task
Implement gemv_kernel and run(M=128, K=256) comparing to A @ x.
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:
M = 128, K = 256
True
- The input values
M = 128andK = 256define the dimensions of matrixAand vectorx, whereAis a (M,K) matrix andxis a length-K vector. - The
gemv_kernelfunction is implemented to compute the matrix-vector producty = A @ xusing a reduction, where each program loads a row ofAand the full vectorx, multiplies elementwise, and reduces withtl.sum. - The
runfunction allocates inputs on the GPU, launches thegemv_kernel, and computes the reference result usingA @ x. - The final output is
Trueif the result from thegemv_kernelis close to the reference result, i.e.,torch.allclose(triton_out, torch_reference, ...)returnsTrue.
Constraints:
- One program per row, grid = (M,)
- BLOCK_K = triton.next_power_of_2(K); mask columns
- y[m] = tl.sum(row * x, axis=0)
Background Knowledge
The problem involves computing the matrix-vector product, a fundamental operation in linear algebra. Given a matrix A of size (M, K) and a vector x of length K, the goal is to compute the product y = A @ x. This operation is a key component of many machine learning and scientific computing applications. The matrix-vector product is defined as the sum of the products of the elements of each row of A with the corresponding elements of x. In other words, the i-th element of the output vector y is computed as the dot product of the i-th row of A and the vector x.
The problem also mentions Triton, a programming language and framework for writing high-performance, GPU-accelerated code. Triton provides a Python-like interface for defining kernels, which are small programs that run on the GPU. In this problem, we need to define a Triton kernel that computes the matrix-vector product for a single row of the output vector y. The kernel will load the corresponding row of A and the entire vector x, perform an element-wise multiplication, and then reduce the result using tl.sum.
The problem also highlights the importance of choosing the right approach for computing the matrix-vector product. In this case, we are asked to use a reduction approach, where each program computes a single row of the output vector y. This approach can be more efficient than using a matrix multiplication (matmul) when the inner dimension K is small.
Algorithm/Approach
The general approach to solving this problem involves defining a Triton kernel that computes the matrix-vector product for a single row of the output vector y. The kernel will perform the following operations:
- Load the corresponding row of A and the entire vector x
- Perform an element-wise multiplication of the row of A and the vector x
- Reduce the result using tl.sum to compute the corresponding element of the output vector y
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.