Implement a fused linear-layer epilogue in Triton: matmul, add a per-column bias, apply ReLU, all in one kernel. Given A of shape (M,K), B of shape (K,N), and bias of shape (N,):
out[i,j]=max(0,k=0∑K−1A[i,k]⋅B[k,j]+bias[j])The epilogue is applied in registers on the accumulator tile, before the store. This avoids the round-trip to global memory that an unfused implementation would incur, and is the simplest concrete example of why fused kernels dominate transformer inference workloads.
Input: A = [[1, 2], [3, 4]], B = [[1, 0], [0, 1]], bias = [0, -10]
Output: [[1, 0], [3, 0]]
Input: A = [[1, -1], [-1, 1]], B = [[1, 1], [1, 1]], bias = [1, 1]
Output: [[1, 1], [1, 1]]
Start from the plain tiled matmul template. The fusion lives in the few lines between the K-loop and the final store.
Load the bias vector once per program for the columns of the output tile and broadcast it across the row axis when adding to the accumulator.
ReLU is a branch-free elementwise max against zero on the accumulator. Apply it after the bias add and before the store.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for K≥128.Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Implement a fused linear-layer epilogue in Triton: matmul, add a per-column bias, apply ReLU, all in one kernel. Given A of shape (M,K), B of shape (K,N), and bias of shape (N,):
out[i,j]=max(0,k=0∑K−1A[i,k]⋅B[k,j]+bias[j])The epilogue is applied in registers on the accumulator tile, before the store. This avoids the round-trip to global memory that an unfused implementation would incur, and is the simplest concrete example of why fused kernels dominate transformer inference workloads.
Input: A = [[1, 2], [3, 4]], B = [[1, 0], [0, 1]], bias = [0, -10]
Output: [[1, 0], [3, 0]]
Input: A = [[1, -1], [-1, 1]], B = [[1, 1], [1, 1]], bias = [1, 1]
Output: [[1, 1], [1, 1]]
Start from the plain tiled matmul template. The fusion lives in the few lines between the K-loop and the final store.
Load the bias vector once per program for the columns of the output tile and broadcast it across the row axis when adding to the accumulator.
ReLU is a branch-free elementwise max against zero on the accumulator. Apply it after the bias add and before the store.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for K≥128.Accepts: array
Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.