1 / 1

Tiled Matrix Multiplication in Triton

Linear AlgebraGPU Programming

Medium

Implement a tiled matrix multiplication kernel in Triton. Given $A$ of shape $(M, K)$ and $B$ of shape $(K, N)$ , compute $C = A \cdot B$ of shape $(M, N)$ :

C[i, j] = \sum_{k=0}^{K-1} A[i, k] \cdot B[k, j]

This is the canonical 2D-tile parallel pattern. Each program owns a $\texttt{BLOCK\_M} \times \texttt{BLOCK\_N}$ output tile and walks the shared $K$ dimension in chunks, accumulating partial products in an fp32 register tile before storing the result.

Loading visualization...

Examples

Input: A = [[1, 2], [3, 4]], B = [[5, 6], [7, 8]]

Output: [[19, 22], [43, 50]]

Input: A = [[1, 0, 1], [0, 1, 0]], B = [[1, 2], [3, 4], [5, 6]]

Output: [[6, 8], [3, 4]]

Hint 1

Use a 2D launch grid sized $(\lceil M / \text{BLOCK\_M} \rceil, \lceil N / \text{BLOCK\_N} \rceil)$ . Each program reads its two program IDs to locate its output tile.

Hint 2

Initialize the accumulator with $\texttt{tl.zeros}((\text{BLOCK\_M}, \text{BLOCK\_N}), \text{dtype}=\texttt{tl.float32})$ , loop over $K$ in steps of $\text{BLOCK\_K}$ , and cast to the output dtype only on the final store.

Hint 3

Use $\texttt{tl.dot}$ on the loaded $A$ and $B$ tiles. The mask on the $K$ tail prevents partial products from leaking past the last full chunk.

Requirements

Use one Triton program per output tile on a 2D grid
Walk the $K$ dimension in tile-sized chunks, accumulating in fp32
Mask the $M$ , $N$ , and $K$ tails so non-power-of-two shapes are correct
The $\texttt{solve}$ signature must not change
use ieee precision for consistent float32 accuracy across gpus.

Constraints

$1 \le M, N, K \le 512$ .
All tensors are torch.float32 on cuda.
Tolerance: atol=1e-2 + rtol=1e-2 * abs(expected) for $K \ge 128$ to absorb fp32 accumulation drift.

Try Similar Problems

Autotuned Matrix Multiplication in TritonMedium Tiled Transpose in TritonMedium GEMV: Matrix Vector Product in TritonMedium Fused Matmul + Bias + ReLU in TritonMedium Scaled Dot-Product Attention in TritonMedium Block-Pointer Matmul in TritonMedium Split-K Matmul in TritonHard Vector Addition in TritonEasy Fused Multiply-Add in TritonEasy Grouped Program-ID Matmul in TritonHard L2 Vector Norm in TritonEasy RMS Normalization as a CUDA KernelMedium

Tiled Matrix Multiplication in Triton

Linear AlgebraGPU Programming

Medium

Implement a tiled matrix multiplication kernel in Triton. Given $A$ of shape $(M, K)$ and $B$ of shape $(K, N)$ , compute $C = A \cdot B$ of shape $(M, N)$ :

C[i, j] = \sum_{k=0}^{K-1} A[i, k] \cdot B[k, j]

Loading visualization...

Examples

Input: A = [[1, 2], [3, 4]], B = [[5, 6], [7, 8]]

Output: [[19, 22], [43, 50]]

Input: A = [[1, 0, 1], [0, 1, 0]], B = [[1, 2], [3, 4], [5, 6]]

Output: [[6, 8], [3, 4]]

Hint 1

Use a 2D launch grid sized $(\lceil M / \text{BLOCK\_M} \rceil, \lceil N / \text{BLOCK\_N} \rceil)$ . Each program reads its two program IDs to locate its output tile.

Hint 2

Hint 3

Use $\texttt{tl.dot}$ on the loaded $A$ and $B$ tiles. The mask on the $K$ tail prevents partial products from leaking past the last full chunk.

Requirements

Use one Triton program per output tile on a 2D grid
Walk the $K$ dimension in tile-sized chunks, accumulating in fp32
Mask the $M$ , $N$ , and $K$ tails so non-power-of-two shapes are correct
The $\texttt{solve}$ signature must not change
use ieee precision for consistent float32 accuracy across gpus.

Constraints

$1 \le M, N, K \le 512$ .
All tensors are torch.float32 on cuda.
Tolerance: atol=1e-2 + rtol=1e-2 * abs(expected) for $K \ge 128$ to absorb fp32 accumulation drift.

Try Similar Problems

A =

Accepts: array

B =

Accepts: array

Examples

Requirements

Constraints

Kernel performance

Tiled Matrix Multiplication in Triton

Examples

Requirements

Constraints

Kernel performance