Function blockwise_bmm

tests/python/contrib/test_cutlass_gemm.py:275–302

(
    x_fp8_np: np.ndarray,
    x_scale_np: np.ndarray,
    w_np: np.ndarray,
    w_scale_np: np.ndarray,
    block_size: tuple[int, int],
    dtype: str,
)

def blockwise_bmm(
    x_fp8_np: np.ndarray,
    x_scale_np: np.ndarray,
    w_np: np.ndarray,
    w_scale_np: np.ndarray,
    block_size: tuple[int, int],
    dtype: str,
):
    o_np = np.zeros((x_fp8_np.shape[0], x_fp8_np.shape[1], w_np.shape[1]), dtype=dtype)
    for j in range(w_scale_np.shape[1]):
        for k in range(w_scale_np.shape[2]):
            o_np[..., j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[1])] += (
                np.matmul(
                    x_fp8_np[
                        ..., k * block_size[1] : min((k + 1) * block_size[1], x_fp8_np.shape[2])
                    ].astype(dtype),
                    w_np[
                        ...,
                        j * block_size[0] : min((j + 1) * block_size[0], w_np.shape[1]),
                        k * block_size[1] : min((k + 1) * block_size[1], w_np.shape[2]),
                    ]
                    .transpose(0, 2, 1)
                    .astype(dtype),
                )
                * x_scale_np[..., k : k + 1]
                * w_scale_np[..., j : j + 1, k : k + 1]
            )
    return o_np


@pytest.mark.skipif(not env.has_cutlass(), reason="need cutlass")