Triton Matmul with Bias Epilogue
Problem Statement
Extend tiled matmul to add a per-column bias vector after accumulation: C = A @ B + bias, where bias has length N.
Background
The bias add is an "epilogue" applied to the accumulator before storing. Load bias[offs_n] and broadcast it across the tile rows.
Your Task
Implement matmul_bias_kernel and run(M=128, N=128, K=128) comparing to A @ B + bias.
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 input matricesAandBof size M×K and K×N respectively, and a bias vectorbiasof length N, on the GPU. - It then launches the
matmul_bias_kernelfunction, which performs a tiled matrix multiplication ofAandBand adds the bias vector to each column of the result. - The result of the
matmul_bias_kernelfunction is stored intriton_outand compared to the reference resulttorch_reference, calculated asA @ B + bias. - The comparison is done using
torch.allclose, which checks if the two tensors are element-wise equal within a certain tolerance, and returnsTrueif they are equal, indicating that the implementation is correct.
Constraints:
- Reuse the tiled-matmul structure
- Epilogue: acc += bias loaded at offs_n (broadcast over rows)
- Mask bias load with offs_n < N
Background Knowledge
The problem involves matrix multiplication with an added bias vector, which is a common operation in deep learning. Matrix multiplication is a fundamental concept in linear algebra, where two matrices A and B are multiplied to produce a new matrix C. The operation is defined as C = A @ B, where @ denotes the matrix product. In this case, we need to add a bias vector to the result of the matrix multiplication, which is a common technique used in neural networks to introduce a shift in the output.
The Triton programming framework is used to implement the matrix multiplication with bias. Triton is a Python-based programming language and framework for writing high-performance GPU code. It allows developers to write custom kernels that can be executed on NVIDIA GPUs. The tiled matmul approach is used to divide the matrix into smaller tiles, which can be processed in parallel on the GPU. This approach is useful for large matrices that do not fit in the GPU's memory.
The epilogue is a term used to describe the additional operations that are applied to the accumulator before storing the result. In this case, the epilogue is the addition of the bias vector to the result of the matrix multiplication. The bias vector is loaded and broadcasted across the tile rows, which means that each element of the bias vector is added to the corresponding element of the tile.
Algorithm/Approach
The general approach to solving this problem involves the following steps:
- Divide the matrix into smaller tiles that can be processed in parallel on the GPU
- Implement a custom kernel that performs the matrix multiplication and adds the bias vector to the result
- Use the Triton programming framework to launch the kernel on the GPU and store the result
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.