| 248 | |
| 249 | |
| 250 | def matmul_tma(a, b, warp_specialize: bool): |
| 251 | # Check constraints. |
| 252 | assert a.shape[1] == b.shape[1], "Incompatible dimensions" # b is transposed |
| 253 | assert a.dtype == b.dtype, "Incompatible dtypes" |
| 254 | |
| 255 | M, K = a.shape |
| 256 | N, K = b.shape |
| 257 | dtype = a.dtype |
| 258 | |
| 259 | c = torch.empty((M, N), device=a.device, dtype=dtype) |
| 260 | |
| 261 | # A dummy block value that will be overwritten when we have the real block size |
| 262 | dummy_block = [1, 1] |
| 263 | a_desc = TensorDescriptor.from_tensor(a, dummy_block) |
| 264 | b_desc = TensorDescriptor.from_tensor(b, dummy_block) |
| 265 | c_desc = TensorDescriptor.from_tensor(c, dummy_block) |
| 266 | |
| 267 | def grid(META): |
| 268 | BLOCK_M = META["BLOCK_SIZE_M"] |
| 269 | BLOCK_N = META["BLOCK_SIZE_N"] |
| 270 | return (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),) |
| 271 | |
| 272 | matmul_kernel_tma[grid]( |
| 273 | a_desc, |
| 274 | b_desc, |
| 275 | c_desc, # |
| 276 | M, |
| 277 | N, |
| 278 | K, # |
| 279 | FP8_OUTPUT=dtype == torch.float8_e4m3fn, # |
| 280 | WARP_SPECIALIZE=warp_specialize, # |
| 281 | ) |
| 282 | return c |
| 283 | |
| 284 | |
| 285 | @triton.jit |