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
Input
M = 128, N = 128, K = 128
Output
True
Reasoning
- The
runfunction allocates input matricesAandBof size and respectively, and a bias vectorbiasof length , 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