TensorTonic
Study Plans
Study Plan
Fused Tiled Matmul and ReLU
Activation FunctionsLinear AlgebraGPU Programming
Hard

Implement fused matrix multiplication and ReLU in Triton. Input matrices AAA and BBB have shapes (M,K)(M,K)(M,K) and (K,N)(K,N)(K,N), and the output YYY has shape (M,N)(M,N)(M,N):

Ym,n=max⁡(0,∑k=0K−1float32⁡(Am,k)float32⁡(Bk,n)).Y_{m,n}=\max\left(0,\sum_{k=0}^{K-1}\operatorname{float32}(A_{m,k})\operatorname{float32}(B_{k,n})\right).Ym,n​=max(0,k=0∑K−1​float32(Am,k​)float32(Bk,n​)).

Here, MMM and NNN are output dimensions, KKK is the shared inner dimension, and mmm, nnn, and kkk index those dimensions. Float32 denotes the arithmetic precision, and the maximum applies ReLU. Complete only the empty kernel body; the supplied solve launcher is implemented. The pointer arguments identify the three buffers, the size arguments give M, N, and K, the stride arguments give each buffer's element strides, and BLOCK_M, BLOCK_N, and BLOCK_K are compile-time tile sizes. Return no value; write Y into the supplied output tensor using its preallocated dtype.

Examples

Input: a = [[1, -2], [3, 1]], b = [[2, -1, 0], [1, 2, -3]], dtype = float16, output_dtype = float16

Output: [[0, 0, 6], [7, 0, 0]]

Explanation: ReLU replaces the negative entries of the matrix product with zero.

Input: a = [[1, 2, 3], [0, -1, 2]], b = [[1, 0], [0, 1], [1, -1]], dtype = bfloat16, output_dtype = bfloat16

Output: [[4, 0], [2, 0]]

Input: a = [[-1, 2], [2, -1]], b = [[1, -3], [4, 2]], dtype = float16, output_dtype = float32

Output: [[7, 7], [0, 0]]

Hint 1

tl.program_id can address each dimension of a program grid.

Hint 2

tl.dot supports float32 accumulation of tile products.

Hint 3

tl.load and tl.store support masks for partial tiles.

Requirements

  • Use Triton with float32 accumulation for the matrix product.
  • Apply ReLU to the complete matrix-product values before storage conversion.
  • Respect the supplied strides and preserve both input tensors.
  • A zero inner dimension produces zeros; zero output rows or columns require no writes.

Constraints

  • a and b are two-dimensional CUDA tensors with matching float16 or bfloat16 storage dtype and equal inner dimensions.
  • out is a separate tensor on the same CUDA device, shaped (M, N), with float16, bfloat16, or float32 storage dtype.
  • Input values are finite; nonempty input and output tensors have positive strides.
  • The supplied launcher handles empty dimensions and rejects incompatible shapes or unsupported dtypes.
  • The supplied launcher converts bfloat16 inputs to float32 before kernel execution.
Try Similar Problems
Triton Row-Wise SoftmaxMediumFused Matmul + Bias + ReLU in TritonMediumMasked Triton GELU KernelEasyBlockwise Online SoftmaxHardSplit-K Matmul in TritonHardTiled Matrix Multiplication as a CUDA KernelHardGrouped Program-ID Matmul in TritonHardScaled Dot-Product Attention in CUDAHardTiled Triton Row SumMediumRMS Normalization as a CUDA KernelMediumScaled Dot-Product Attention in TritonMediumBlock-Pointer Matmul in TritonMedium

Accepts: array

Accepts: array

Accepts: string

Accepts: string

You must run your code first.

Kernel performance

Submit a correct solution to see your kernel performance.

Fused Tiled Matmul and ReLU

Activation FunctionsLinear AlgebraGPU Programming
Hard

Implement fused matrix multiplication and ReLU in Triton. Input matrices AAA and BBB have shapes (M,K)(M,K)(M,K) and (K,N)(K,N)(K,N), and the output YYY has shape (M,N)(M,N)(M,N):

Ym,n=max⁡(0,∑k=0K−1float32⁡(Am,k)float32⁡(Bk,n)).Y_{m,n}=\max\left(0,\sum_{k=0}^{K-1}\operatorname{float32}(A_{m,k})\operatorname{float32}(B_{k,n})\right).Ym,n​=max(0,k=0∑K−1​float32(Am,k​)float32(Bk,n​)).

Here, MMM and NNN are output dimensions, KKK is the shared inner dimension, and mmm, nnn, and kkk index those dimensions. Float32 denotes the arithmetic precision, and the maximum applies ReLU. Complete only the empty kernel body; the supplied solve launcher is implemented. The pointer arguments identify the three buffers, the size arguments give M, N, and K, the stride arguments give each buffer's element strides, and BLOCK_M, BLOCK_N, and BLOCK_K are compile-time tile sizes. Return no value; write Y into the supplied output tensor using its preallocated dtype.

Examples

Input: a = [[1, -2], [3, 1]], b = [[2, -1, 0], [1, 2, -3]], dtype = float16, output_dtype = float16

Output: [[0, 0, 6], [7, 0, 0]]

Explanation: ReLU replaces the negative entries of the matrix product with zero.

Input: a = [[1, 2, 3], [0, -1, 2]], b = [[1, 0], [0, 1], [1, -1]], dtype = bfloat16, output_dtype = bfloat16

Output: [[4, 0], [2, 0]]

Input: a = [[-1, 2], [2, -1]], b = [[1, -3], [4, 2]], dtype = float16, output_dtype = float32

Output: [[7, 7], [0, 0]]

Hint 1

tl.program_id can address each dimension of a program grid.

Hint 2

tl.dot supports float32 accumulation of tile products.

Hint 3

tl.load and tl.store support masks for partial tiles.

Requirements

  • Use Triton with float32 accumulation for the matrix product.
  • Apply ReLU to the complete matrix-product values before storage conversion.
  • Respect the supplied strides and preserve both input tensors.
  • A zero inner dimension produces zeros; zero output rows or columns require no writes.

Constraints

  • a and b are two-dimensional CUDA tensors with matching float16 or bfloat16 storage dtype and equal inner dimensions.
  • out is a separate tensor on the same CUDA device, shaped (M, N), with float16, bfloat16, or float32 storage dtype.
  • Input values are finite; nonempty input and output tensors have positive strides.
  • The supplied launcher handles empty dimensions and rejects incompatible shapes or unsupported dtypes.
  • The supplied launcher converts bfloat16 inputs to float32 before kernel execution.
Try Similar Problems
Triton Row-Wise SoftmaxMediumFused Matmul + Bias + ReLU in TritonMediumMasked Triton GELU KernelEasyBlockwise Online SoftmaxHardSplit-K Matmul in TritonHardTiled Matrix Multiplication as a CUDA KernelHardGrouped Program-ID Matmul in TritonHardScaled Dot-Product Attention in CUDAHardTiled Triton Row SumMediumRMS Normalization as a CUDA KernelMediumScaled Dot-Product Attention in TritonMediumBlock-Pointer Matmul in TritonMedium

Accepts: array

Accepts: array

Accepts: string

Accepts: string

You must run your code first.

Kernel performance

Submit a correct solution to see your kernel performance.