Triton Matmul with ReLU Epilogue
Problem Statement
Fuse a ReLU into the matmul epilogue: C = relu(A @ B).
Background
Apply tl.maximum(acc, 0.0) to the accumulator before storing. Fusing the activation avoids a second kernel reading and writing all of C.
Your Task
Implement matmul_relu_kernel and run(M=128, N=128, K=128) comparing to torch.relu(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
runfunction allocates two input matricesAandBof size M×K and K×N respectively on the GPU, where M=128, N=128, and K=128. - It then launches the
matmul_relu_kernelfunction, which performs the matrix multiplicationA @ Band applies the ReLU activation function to the result, effectively computing C=max(A@B,0.0). - The result from the
matmul_relu_kernelfunction is stored intriton_outand compared to the result fromtorch.relu(A @ B)usingtorch.allclose. - The comparison returns
Trueif the two results are close within a certain tolerance, indicating that the custom implementation matches the PyTorch reference implementation.
Constraints:
- Tiled matmul accumulator
- Epilogue: acc = tl.maximum(acc, 0.0) before store
Background Knowledge
The problem involves matrix multiplication and the application of the ReLU (Rectified Linear Unit) activation function. Matrix multiplication is a fundamental operation in linear algebra, where two matrices A and B are multiplied to produce another matrix C. The ReLU function, on the other hand, is a widely used activation function in deep learning, which maps all negative values to 0 and all positive values to the same value. In this problem, we need to fuse the ReLU function into the matrix multiplication operation, which means applying the ReLU function to the result of the matrix multiplication before storing it.
The Triton programming language is a Python-based language for writing high-performance GPU code. It provides a simple and intuitive way to write custom GPU kernels, which can be used to accelerate various computations, including matrix multiplication. In this problem, we need to implement a Triton kernel that performs matrix multiplication with ReLU activation. The kernel will take two input matrices A and B and produce an output matrix C, where C = relu(A @ B).
To solve this problem, we need to understand the basics of GPU programming and parallel computing. We need to know how to launch a GPU kernel, how to manage memory on the GPU, and how to optimize the performance of the kernel. We also need to understand the Triton programming model, which provides a set of APIs and tools for writing and optimizing GPU code.
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.