Function get_ref_data

tests/python/contrib/test_cutlass_gemm.py:45–55

()

    @memoize("tvm.contrib.cutlass.test_group_gemm_sm90")
    def get_ref_data():
        assert M % num_groups == 0
        M_per_group = M // num_groups
        a_np = get_random_tensor((M, K), x_dtype)
        b_np = get_random_tensor((num_groups, N, K), weight_dtype)
        indptr_np = np.arange(1, num_groups + 1).astype("int64") * M_per_group
        c_np = np.concatenate(
            [a_np[i * M_per_group : (i + 1) * M_per_group] @ b_np[i].T for i in range(num_groups)],
            axis=0,
        )
        return a_np, b_np, indptr_np, c_np

    def to_numpy_dtype(dtype):
        mapping = {"float8_e5m2": ml_dtypes.float8_e5m2, "float8_e4m3fn": ml_dtypes.float8_e4m3fn}