Triton Matrix-Vector Product (GEMV)

mediumPROMatrix MultiplicationTriton Programming

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

Input

M = 128, K = 256

Output

True

Reasoning

The input values M = 128 and K = 256 define the dimensions of matrix A and vector x, where A is a $(M, K)$ matrix and x is a length- $K$ vector.
The gemv_kernel function is implemented to compute the matrix-vector product y = A @ x using a reduction, where each program loads a row of A and the full vector x, multiplies elementwise, and reduces with tl.sum.
The run function allocates inputs on the GPU, launches the gemv_kernel, and computes the reference result using A @ x.
The final output is True if the result from the gemv_kernel is close to the reference result, i.e., torch.allclose(triton_out, torch_reference, ...) returns True.

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)

Learn this in

Why GPUs? The Parallel Computing MindsetGPU Basics with Triton Matrix MultiplicationGPU Basics with Triton

More from Triton Programming

Triton Single-Query Attention Kernelhard Triton Row Cumulative Sumhard Triton Stable LogSumExp Kernelhard Triton Seeded Fused Dropouthard Triton Autotuned Vector Addmedium

Concept primerPREVIEW

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

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.

PRO

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.

solution.py

Python 3.13

GPU · T4

Test results