Triton Tiled Matrix Multiplication
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:
M = 128, N = 128, K = 128
True
- The input values
M = 128,N = 128, andK = 128are used to allocate matricesAandBwith shapes(M, K)and(K, N)respectively. - The
matmul_kernelfunction is launched with a 2D grid to compute the matrix productC = A @ Bin a tiled manner, using blocks of sizeBLOCK_M x BLOCK_Nand accumulating results in a float32 register tile. - The tiled matrix multiplication is performed by looping over
Kin steps ofBLOCK_K, loading tiles ofAandB, and multiplying them usingtl.dot, with masks handling non-tile-multiple sizes. - The resulting matrix
Cfrom the Triton kernel is compared to the reference result fromtorchusingtorch.allclose, which checks if the two matrices are element-wise equal within a certain tolerance, resulting in the outputTrue.
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
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⋅N⋅K/(BLOCKM⋅BLOCKN⋅BLOCKK)), 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⋅N+K⋅N+M⋅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.