Triton Scaled Matmul
Problem Statement
Compute a scaled matmul C = alpha * (A @ B) with alpha a runtime scalar.
Background
Multiply the accumulator by alpha in the epilogue. This mirrors BLAS GEMM's alpha parameter.
Your Task
Implement scaled_matmul_kernel and run(M=128, N=128, K=128, alpha=0.5) comparing to alpha * (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, alpha = 0.5
True
- We start with the given input values: M = 128, N = 128, K = 128, and alpha = 0.5.
- The
scaled_matmul_kernelfunction computes the matrix product of A and B, resulting in a temporary matrix product A@B. - The kernel then scales this temporary result by the given alpha value, 0.5⋅(A@B), to produce the final output matrix C.
- The output of the
runfunction isTrueif the resulting matrix C is close to the reference solution computed byalpha * (A @ B), which is verified usingtorch.allclose.
Constraints:
- Tiled matmul accumulator
- Epilogue: acc = acc * alpha
- alpha is a runtime scalar
Background Knowledge
Introduction to Matrix Multiplication
Matrix multiplication is a fundamental operation in linear algebra, where two matrices A and B are multiplied to produce another matrix C. The resulting matrix C has dimensions that are a combination of the dimensions of A and B. In this problem, we are dealing with a scaled matrix multiplication, where the result of the multiplication is scaled by a runtime scalar alpha.
Understanding Scaled Matmul
The scaled matmul operation can be represented as C = alpha * (A @ B), where @ denotes the matrix multiplication operation. This operation is commonly used in various fields such as machine learning, scientific computing, and data analysis. The alpha parameter allows for scaling the result of the matrix multiplication, which can be useful in various applications.
Triton Programming and GPU Acceleration
Triton is a programming language and framework that allows for writing high-performance, GPU-accelerated code. In this problem, we are required to implement the scaled matmul operation using Triton, which involves defining a kernel function that can be executed on the GPU. The use of GPU acceleration can significantly improve the performance of matrix multiplication operations, especially for large matrices.
Algorithm/Approach
The general approach to solving this problem involves implementing a kernel function in Triton that performs the scaled matmul operation. This can be achieved by using the following algorithm pattern:
- Define the kernel function that takes in the input matrices A and B, the scalar alpha, and the output matrix C.
- Use Triton's programming model to launch the kernel function on the GPU, which involves defining the block and grid dimensions.
- Within the kernel function, perform the matrix multiplication operation and scale the result by alpha.
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.