(
func_name, M, N, K, num_groups, x_dtype, weight_dtype, out_dtype, use_scale, rtol, atol
)
| 34 | |
| 35 | |
| 36 | def verify_group_gemm( |
| 37 | func_name, M, N, K, num_groups, x_dtype, weight_dtype, out_dtype, use_scale, rtol, atol |
| 38 | ): |
| 39 | group_gemm_func = tvm.get_global_func(func_name, allow_missing=True) |
| 40 | if group_gemm_func is None: |
| 41 | print(f"Skipped as {func_name} is not available") |
| 42 | return |
| 43 | |
| 44 | @memoize("tvm.contrib.cutlass.test_group_gemm_sm90") |
| 45 | def get_ref_data(): |
| 46 | assert M % num_groups == 0 |
| 47 | M_per_group = M // num_groups |
| 48 | a_np = get_random_tensor((M, K), x_dtype) |
| 49 | b_np = get_random_tensor((num_groups, N, K), weight_dtype) |
| 50 | indptr_np = np.arange(1, num_groups + 1).astype("int64") * M_per_group |
| 51 | c_np = np.concatenate( |
| 52 | [a_np[i * M_per_group : (i + 1) * M_per_group] @ b_np[i].T for i in range(num_groups)], |
| 53 | axis=0, |
| 54 | ) |
| 55 | return a_np, b_np, indptr_np, c_np |
| 56 | |
| 57 | def to_numpy_dtype(dtype): |
| 58 | mapping = {"float8_e5m2": ml_dtypes.float8_e5m2, "float8_e4m3fn": ml_dtypes.float8_e4m3fn} |
| 59 | return mapping.get(dtype, dtype) |
| 60 | |
| 61 | a_np, b_np, indptr_np, c_np = get_ref_data() |
| 62 | dev = tvm.cuda(0) |
| 63 | a_nd = tvm.runtime.tensor(a_np.astype(to_numpy_dtype(x_dtype)), device=dev) |
| 64 | b_nd = tvm.runtime.tensor(b_np.astype(to_numpy_dtype(weight_dtype)), device=dev) |
| 65 | c_nd = tvm.runtime.empty(c_np.shape, dtype=out_dtype, device=dev) |
| 66 | indptr_nd = tvm.runtime.tensor(indptr_np, device=dev) |
| 67 | workspace = tvm.runtime.empty((4096 * 1024,), dtype="uint8", device=dev) |
| 68 | if use_scale: |
| 69 | scale = tvm.runtime.tensor(np.array([1.0], dtype="float32"), device=dev) |
| 70 | group_gemm_func(a_nd, b_nd, indptr_nd, workspace, scale, c_nd) |
| 71 | else: |
| 72 | group_gemm_func(a_nd, b_nd, indptr_nd, workspace, c_nd) |
| 73 | tvm.testing.assert_allclose(c_nd.numpy(), c_np, rtol=rtol, atol=atol) |
| 74 | |
| 75 | |
| 76 | @pytest.mark.skipif(not env.has_cutlass(), reason="need cutlass") |
no test coverage detected
searching dependent graphs…