Triton Scaled Matmul

hardPROMatrix MultiplicationTriton Programming

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

Input

M = 128, N = 128, K = 128, alpha = 0.5

Output

True

Reasoning

We start with the given input values: M = 128, N = 128, K = 128, and alpha = 0.5.
The scaled_matmul_kernel function 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 \cdot (A @ B)$ , to produce the final output matrix C.
The output of the run function is True if the resulting matrix C is close to the reference solution computed by alpha * (A @ B), which is verified using torch.allclose.

Constraints

Tiled matmul accumulator
Epilogue: acc = acc * alpha
alpha is a runtime scalar

Learn this in

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 primerPREVIEW

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.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

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