cutlass_python_jit_gemm_matrix_multiplication_quickstart.py

1import torch
2import cutlass
3
4# 1. Define the problem size
5M, N, K = 1024, 1024, 1024
6
7# 2. Create input tensors (using PyTorch as the backend)
8# CUTLASS Python supports torch.Tensor and numpy.ndarray
9A = torch.randn((M, K), dtype=torch.float16, device="cuda")
10B = torch.randn((K, N), dtype=torch.float16, device="cuda")
11C = torch.zeros((M, N), dtype=torch.float16, device="cuda")
12
13# 3. Create and configure the GEMM operation
14# The 'cutlass.op.Gemm' interface automatically selects an appropriate kernel
15plan = cutlass.op.Gemm(element=torch.float16, layout=cutlass.LayoutType.RowMajor)
16
17# 4. Run the operation
18# This will JIT-compile the kernel on the first call if not already cached
19plan.run(A, B, C)
20
21# 5. Verify the result
22expected = torch.mm(A, B)
23torch.testing.assert_close(C, expected, atol=1e-2, rtol=1e-2)
24print("CUTLASS GEMM execution successful and verified.")