| 154 | |
| 155 | |
| 156 | def matmul(a, b): |
| 157 | # Check constraints. |
| 158 | assert a.shape[1] == b.shape[0], "Incompatible dimensions" |
| 159 | assert a.dtype == b.dtype, "Incompatible dtypes" |
| 160 | M, K = a.shape |
| 161 | K, N = b.shape |
| 162 | dtype = a.dtype |
| 163 | |
| 164 | c = torch.empty((M, N), device=a.device, dtype=dtype) |
| 165 | # 1D launch kernel where each block gets its own program. |
| 166 | grid = lambda META: ( |
| 167 | triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), |
| 168 | ) |
| 169 | matmul_kernel[grid]( |
| 170 | a, |
| 171 | b, |
| 172 | c, # |
| 173 | M, |
| 174 | N, |
| 175 | K, # |
| 176 | a.stride(0), |
| 177 | a.stride(1), # |
| 178 | b.stride(0), |
| 179 | b.stride(1), # |
| 180 | c.stride(0), |
| 181 | c.stride(1), # |
| 182 | ) |
| 183 | return c |
| 184 | |
| 185 | |
| 186 | def matmul_tma_set_block_size_hook(nargs): |