Triton Matmul with ReLU Epilogue

hardPROMatrix MultiplicationTriton Programming

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

Input

M = 128, N = 128, K = 128

Output

True

Reasoning

The run function allocates two input matrices A and B of size $M \times K$ and $K \times N$ respectively on the GPU, where $M = 128$ , $N = 128$ , and $K = 128$ .
It then launches the matmul_relu_kernel function, which performs the matrix multiplication A @ B and applies the ReLU activation function to the result, effectively computing $C = \max(A @ B, 0.0)$ .
The result from the matmul_relu_kernel function is stored in triton_out and compared to the result from torch.relu(A @ B) using torch.allclose.
The comparison returns True if 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

Learn this in

Memory Access, Autotuning & BenchmarkingGPU Basics with Triton Matrix MultiplicationGPU Basics with Triton Advanced Kernels & OptimizationGPU 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

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