Triton Tiled Matrix Multiplication

hardFREEMatrix MultiplicationTriton Programming

Problem Statement

Implement a tiled matrix multiplication C = A @ B where A is (M, K) and B is (K, N), using a 2D launch grid and an accumulator.

Background

Each program computes one BLOCK_M x BLOCK_N output tile. Loop over K in steps of BLOCK_K, loading tiles of A and B, multiplying with tl.dot, and accumulating into a float32 register tile. Strides let the kernel work on any contiguous layout; masks handle non-tile-multiple sizes.

Your Task

Implement matmul_kernel and run(M=128, N=128, K=128) comparing to A @ B.

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, N = 128, K = 128

Output

True

Reasoning

The input values M = 128, N = 128, and K = 128 are used to allocate matrices A and B with shapes (M, K) and (K, N) respectively.
The matmul_kernel function is launched with a 2D grid to compute the matrix product C = A @ B in a tiled manner, using blocks of size BLOCK_M x BLOCK_N and accumulating results in a float32 register tile.
The tiled matrix multiplication is performed by looping over K in steps of BLOCK_K, loading tiles of A and B, and multiplying them using tl.dot, with masks handling non-tile-multiple sizes.
The resulting matrix C from the Triton kernel is compared to the reference result from torch using torch.allclose, which checks if the two matrices are element-wise equal within a certain tolerance, resulting in the output True.

Constraints

2D grid: (cdiv(M, BLOCK_M), cdiv(N, BLOCK_N))
Accumulate in a tl.float32 register tile, loop over K by BLOCK_K
Use tl.dot and mask all loads/stores; pass strides

Learn this in

Why GPUs? The Parallel Computing MindsetGPU Basics with Triton Introduction to TritonGPU 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 primer

Background Knowledge

Matrix multiplication is a fundamental operation in linear algebra, and it's a crucial component in many machine learning and deep learning algorithms. The problem statement involves performing a matrix multiplication C = A @ B, where A is a matrix of size (M, K) and B is a matrix of size (K, N). The result C will be a matrix of size (M, N). To understand this problem, it's essential to have a solid grasp of matrix operations, including matrix multiplication.

In the context of Triton Programming, the problem requires using a 2D launch grid and an accumulator to perform the matrix multiplication. This involves dividing the matrices into smaller tiles and processing them in parallel. The use of blocks and strides allows the kernel to work on any contiguous layout, while masks handle non-tile-multiple sizes. Understanding how to work with these concepts is crucial to solving the problem.

The problem also involves comparing the result of the Triton kernel with the result of the standard matrix multiplication A @ B using torch.allclose. This requires understanding how to allocate inputs on the GPU, launch the Triton kernel, and perform the comparison. The goal is to ensure that the result of the Triton kernel is close to the result of the standard matrix multiplication.

Algorithm/Approach

The general approach to solving this problem involves using a tiled matrix multiplication algorithm. This algorithm divides the matrices into smaller tiles and processes them in parallel using a 2D launch grid. The accumulator is used to accumulate the results of the multiplication of each tile. The algorithm involves looping over the tiles, loading the tiles of A and B, multiplying them using tl.dot, and accumulating the result into a float32 register tile.

Step-by-Step Strategy

To implement the solution, follow these steps:

Define the matmul_kernel function that takes the tiles of A and B as input and returns the result of the multiplication.
Allocate the inputs A and B on the GPU.
Launch the Triton kernel using a 2D launch grid.
Loop over the tiles of A and B in steps of BLOCK_K.
Load the tiles of A and B and multiply them using tl.dot.
Accumulate the result into a float32 register tile.
Compare the result of the Triton kernel with the result of the standard matrix multiplication A @ B using torch.allclose.

Common Pitfalls

When implementing the solution, watch out for the following common pitfalls:

Incorrectly defining the matmul_kernel function or the launch grid.
Failing to allocate the inputs on the GPU or launching the kernel incorrectly.
Incorrectly looping over the tiles or loading the tiles of A and B.
Failing to accumulate the result correctly or comparing the results incorrectly.

Time & Space Complexity

The expected time complexity of the solution is $O(M \cdot N \cdot K / (BLOCK_M \cdot BLOCK_N \cdot BLOCK_K))$ , where M, N, and K are the sizes of the matrices, and BLOCK_M, BLOCK_N, and BLOCK_K are the sizes of the blocks. The space complexity is $O(M \cdot N + K \cdot N + M \cdot K)$ , which is the space required to store the matrices A, B, and C. However, the actual time and space complexity may vary depending on the specific implementation and the hardware used.

solution.py

Python 3.13

GPU · T4

Test results