()
| 4 | |
| 5 | |
| 6 | def test_matmul(): |
| 7 | device = torch.device("cuda") |
| 8 | torch.manual_seed(0) |
| 9 | M, N, K = 100, 200, 304 # noqa: N806 (K must be multiple of 8 for TMA) |
| 10 | a = torch.randn(M, K, device=device, dtype=torch.bfloat16) |
| 11 | b = torch.randn(N, K, device=device, dtype=torch.bfloat16) |
| 12 | c = matmul(a, b) |
| 13 | assert torch.allclose(c, a @ b.T, atol=0.01) |