Reimplement tiled matmul with a split-K parallelization that adds a third axis to the grid along the reduction dimension. For each output entry:
C[i,j]=k=0∑K−1A[i,k]⋅B[k,j]The standard 2D-grid matmul has one program per output tile. When M and N are small relative to K (a tall-skinny GEMM), this leaves most SMs idle. Split-K fixes the under-utilization by letting each (pid_m,pid_n,pid_k) program compute a partial sum over its K slice and atomic-add it into the shared output tile.
With SPLIT_K=4 on a four-tile output grid, the launch goes from four programs to sixteen, which is the difference between five percent and twenty percent SM occupancy on a modern GPU. This is the standard fix called out in the Triton matmul tutorial for skinny GEMMs.
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]]
Pack the three program IDs into a 1D grid of size ⌈M/BLOCK_M⌉⋅⌈N/BLOCK_N⌉⋅SPLIT_K and decode pid_k, pid_m, pid_n inside the kernel.
Each program walks k over only its slice: from pid_k⋅(K/SPLIT_K) to (pid_k+1)⋅(K/SPLIT_K). Accumulate a partial (BLOCK_M,BLOCK_N) tile.
Replace the final tl.store with \texttt{tl.atomic_add} so partials from every pid_k for the same output tile combine correctly. The launcher must zero out first.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for K≥128. The tolerance is intentionally loose because tl.atomic_add accumulation order across programs is non-deterministic, so partial sums combine in fp32 in a different order than the torch reference.Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.
Reimplement tiled matmul with a split-K parallelization that adds a third axis to the grid along the reduction dimension. For each output entry:
C[i,j]=k=0∑K−1A[i,k]⋅B[k,j]The standard 2D-grid matmul has one program per output tile. When M and N are small relative to K (a tall-skinny GEMM), this leaves most SMs idle. Split-K fixes the under-utilization by letting each (pid_m,pid_n,pid_k) program compute a partial sum over its K slice and atomic-add it into the shared output tile.
With SPLIT_K=4 on a four-tile output grid, the launch goes from four programs to sixteen, which is the difference between five percent and twenty percent SM occupancy on a modern GPU. This is the standard fix called out in the Triton matmul tutorial for skinny GEMMs.
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]]
Pack the three program IDs into a 1D grid of size ⌈M/BLOCK_M⌉⋅⌈N/BLOCK_N⌉⋅SPLIT_K and decode pid_k, pid_m, pid_n inside the kernel.
Each program walks k over only its slice: from pid_k⋅(K/SPLIT_K) to (pid_k+1)⋅(K/SPLIT_K). Accumulate a partial (BLOCK_M,BLOCK_N) tile.
Replace the final tl.store with \texttt{tl.atomic_add} so partials from every pid_k for the same output tile combine correctly. The launcher must zero out first.
torch.float32 on cuda.atol=1e-2 + rtol=1e-2 * abs(expected) for K≥128. The tolerance is intentionally loose because tl.atomic_add accumulation order across programs is non-deterministic, so partial sums combine in fp32 in a different order than the torch reference.Accepts: array
Accepts: array
Submit a correct solution to see your kernel performance.